gemini version

.gitignore

@@ -1 +1,6 @@
-Models/multi_phoaudio_gemini/*.pth
+Models/*/*.pth
+data/
+explore/
+.env
+__pycache__/
+*/__pycache__/

Models/gemini_vi/config_gemini_vi_en.yml

@@ -0,0 +1,125 @@
+log_dir: "Models/phoaudio/combine_gemini_vi_en"
+first_stage_path: ""
+save_freq: 1
+log_interval: 50
+device: "cuda"
+epochs_1st: 200 # number of epochs for first stage training (pre-training)
+epochs_2nd: 150 # number of epochs for second stage training (joint training)
+batch_size: 8
+max_len: 400 # maximum number of frames
+# max_len: 800
+pretrained_model: "Models/phoaudio/combine_gemini_vi_en/epoch_2nd_00029.pth"
+second_stage_load_pretrained: true # set to true if the pre-trained model is for 2nd stage
+load_only_params: false # set to true if do not want to load epoch numbers and optimizer parameters
+
+F0_path: "Utils_extend_v1/JDC/bst.t7"
+ASR_config: "Utils_extend_v1/ASR/config.yml"
+ASR_path: "Utils_extend_v1/ASR/epoch_extend_186.pth"
+PLBERT_dir: 'Utils_extend_v1/PLBERT/'
+extend_PLBERT: true # set to true if want to extend the PLBERT model
+
+data_params:
+  train_data: "/home/xdep/data/jupyterhub/users/datnvt/data/custom_datasets/text_gemini_vi_en/train_list.txt"
+  val_data: "/home/xdep/data/jupyterhub/users/datnvt/data/custom_datasets/text_gemini_vi_en/validation_vi_list.txt"
+  root_path: "/home/xdep/data/jupyterhub/users/datnvt/data/custom_datasets/wavs_gemini_phoaudio_multi_speaker_small_v1"
+  OOD_data: "/home/xdep/data/jupyterhub/users/datnvt/data/custom_datasets/text_gemini_phoaudio_multi_speaker_small_v1/ood_multi_phoaudio.txt"
+  min_length: 50 # sample until texts with this size are obtained for OOD texts
+
+symbol: #Total 178 symbols
+  pad: "$"
+  punctuation: ';:,.!?¡¿—…"«»“” '
+  letters: "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+  letters_ipa: "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+  extend: "-124567̪" #ADD MORE SYMBOLS HERE
+
+preprocess_params:
+  sr: 24000
+  spect_params:
+    n_fft: 2048
+    win_length: 1200
+    hop_length: 300
+
+model_params:
+  multispeaker: true
+
+  dim_in: 64 
+  hidden_dim: 512
+  max_conv_dim: 512
+  n_layer: 3
+  n_mels: 80
+
+  n_token: 186 # number of phoneme tokens
+  max_dur: 50 # maximum duration of a single phoneme
+  style_dim: 128 # style vector size
+  
+  dropout: 0.2
+
+  # config for decoder
+  decoder: 
+      type: 'istftnet' # either hifigan or istftnet
+      resblock_kernel_sizes: [3,7,11]
+      upsample_rates :  [10, 6]
+      upsample_initial_channel: 512
+      resblock_dilation_sizes: [[1,3,5], [1,3,5], [1,3,5]]
+      upsample_kernel_sizes: [20, 12]
+      gen_istft_n_fft: 20
+      gen_istft_hop_size: 5
+      
+  # speech language model config
+  slm:
+      model: 'microsoft/wavlm-base-plus'
+      sr: 16000 # sampling rate of SLM
+      hidden: 768 # hidden size of SLM
+      nlayers: 13 # number of layers of SLM
+      initial_channel: 64 # initial channels of SLM discriminator head
+  
+  # style diffusion model config
+  diffusion:
+    embedding_mask_proba: 0.1
+    # transformer config
+    transformer:
+      num_layers: 3
+      num_heads: 8
+      head_features: 64
+      multiplier: 2
+
+    # diffusion distribution config
+    dist:
+      sigma_data: 0.2 # placeholder for estimate_sigma_data set to false
+      estimate_sigma_data: true # estimate sigma_data from the current batch if set to true
+      mean: -3.0
+      std: 1.0
+  
+loss_params:
+    lambda_mel: 5. # mel reconstruction loss
+    lambda_gen: 1. # generator loss
+    lambda_slm: 1. # slm feature matching loss
+    
+    lambda_mono: 1. # monotonic alignment loss (1st stage, TMA)
+    lambda_s2s: 1. # sequence-to-sequence loss (1st stage, TMA)
+    TMA_epoch: 10 # TMA starting epoch (1st stage)
+
+    lambda_F0: 1. # F0 reconstruction loss (2nd stage)
+    lambda_norm: 1. # norm reconstruction loss (2nd stage)
+    lambda_dur: 1. # duration loss (2nd stage)
+    lambda_ce: 20. # duration predictor probability output CE loss (2nd stage)
+    lambda_sty: 1. # style reconstruction loss (2nd stage)
+    lambda_diff: 1. # score matching loss (2nd stage)
+    
+    diff_epoch: 5 # style diffusion starting epoch (2nd stage)
+    joint_epoch: 10 # joint training starting epoch (2nd stage)
+
+optimizer_params:
+  lr: 0.0001 # general learning rate
+  bert_lr: 0.00001 # learning rate for PLBERT
+  ft_lr: 0.00001 # learning rate for acoustic modules
+  
+slmadv_params:
+  min_len: 400 # minimum length of samples
+  max_len: 500 # maximum length of samples
+  batch_percentage: 0.5 # to prevent out of memory, only use half of the original batch size
+  iter: 10 # update the discriminator every this iterations of generator update
+  thresh: 5 # gradient norm above which the gradient is scaled
+  scale: 0.01 # gradient scaling factor for predictors from SLM discriminators
+  sig: 1.5 # sigma for differentiable duration modeling
+  

Models/styles_speaker_parallel.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e92865ecb487a924469f82fdbfdb9dad41cdbd9a58c866d71a323265686ee13
+size 2091581

app.py

@@ -1,40 +1,30 @@
 # -*- coding: utf-8 -*-
 """
-Gradio app.py - wired to your 'inference_one' implementation
-- Reference voice: upload OR choose from train_ref/
-- Uses phonemize_text(), compute_style(), inference_one() exactly like your snippet
+Gradio app.py - StyleTTS2-vi with precomputed style embeddings (.pth)
+- UI có alpha/beta/metrics
+- Style Mixer: 4 slot cố định (Kore, Puck, Algenib, Leda), chỉ chỉnh weight; auto-normalize
+- Luôn hiển thị 4 reference samples (accordion)
+- Không còn dropdown speaker & reference sample auto
 """
 
-import os
-import time
-import yaml
-import numpy as np
-import torch
-import torchaudio
-import librosa
-import gradio as gr
-
+import os, re, glob, time, yaml, torch, librosa, numpy as np, gradio as gr
 from munch import Munch
+from soe_vinorm import SoeNormalizer
 
-# -----------------------------
-# Reproducibility
-# -----------------------------
-torch.manual_seed(0)
-torch.backends.cudnn.benchmark = False
-torch.backends.cudnn.deterministic = True
-np.random.seed(0)
-
-# -----------------------------
-# Device / sample-rate
-# -----------------------------
+# ==============================================================
+# Cấu hình cơ bản
+# ==============================================================
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-SR_OUT = 24000  # target audio rate for synthesis
+SR_OUT = 24000
+ALPHA, BETA, DIFFUSION_STEPS, EMBEDDING_SCALE = 0.0, 0.0, 5, 1.0
 
-# -----------------------------
-# External modules from the project
-# -----------------------------
-from models import *             # noqa: F401,F403
-from utils import *              # noqa: F401,F403
+REF_DIR = "ref_voice"   # thư mục chứa audio mẫu (.wav)
+
+# ==============================================================
+# Import module StyleTTS2
+# ==============================================================
+from models import *
+from utils import *
 from models import build_model
 from text_utils import TextCleaner
 from Utils_extend_v1.PLBERT.util import load_plbert
@@ -42,24 +32,21 @@ from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSche
 
 textcleaner = TextCleaner()
 
-# -----------------------------
-# Config / model loading
-# -----------------------------
+# ==============================================================
+# Load model và checkpoint
+# ==============================================================
 from huggingface_hub import hf_hub_download
-hf_hub_download(
-    repo_id="ltphuongunited/styletts2_vi", 
-    filename="epoch_2nd_00058.pth", 
-    local_dir="Models/multi_phoaudio_gemini", 
-    local_dir_use_symlinks=False,
-)
-
-CONFIG_PATH = os.getenv("MODEL_CONFIG", "Models/multi_phoaudio_gemini/config_phoaudio_gemini_small.yml")
-CHECKPOINT_PTH = os.getenv("MODEL_CKPT", "Models/multi_phoaudio_gemini/epoch_2nd_00058.pth")
-
-# Load config
+# hf_hub_download(
+#     repo_id="ltphuongunited/styletts2_vi", 
+#     filename="gemini_2nd_00045.pth", 
+#     local_dir="Models/gemini_vi", 
+#     local_dir_use_symlinks=False,
+# )
+
+CHECKPOINT_PTH = "Models/gemini_vi/gemini_2nd_00045.pth"
+CONFIG_PATH = "Models/gemini_vi/config_gemini_vi_en.yml"
 config = yaml.safe_load(open(CONFIG_PATH))
 
-# Build components
 ASR_config = config.get("ASR_config", False)
 ASR_path   = config.get("ASR_path", False)
 F0_path    = config.get("F0_path", False)
@@ -68,34 +55,24 @@ PLBERT_dir = config.get("PLBERT_dir", False)
 text_aligner    = load_ASR_models(ASR_path, ASR_config)
 pitch_extractor = load_F0_models(F0_path)
 plbert          = load_plbert(PLBERT_dir)
-
-model_params = recursive_munch(config["model_params"])
+model_params    = recursive_munch(config["model_params"])
 model = build_model(model_params, text_aligner, pitch_extractor, plbert)
 
-# to device & eval
 _ = [model[k].to(DEVICE) for k in model]
 _ = [model[k].eval() for k in model]
 
-# Load checkpoint
-if not os.path.isfile(CHECKPOINT_PTH):
-    raise FileNotFoundError(f"Checkpoint not found at '{CHECKPOINT_PTH}'")
-ckpt = torch.load(CHECKPOINT_PTH, map_location="cpu")
-params = ckpt["net"]
+ckpt = torch.load(CHECKPOINT_PTH, map_location="cpu")["net"]
 for key in model:
-    if key in params:
+    if key in ckpt:
         try:
-            model[key].load_state_dict(params[key])
+            model[key].load_state_dict(ckpt[key])
         except Exception:
             from collections import OrderedDict
-            state_dict = params[key]
             new_state = OrderedDict()
-            for k, v in state_dict.items():
-                name = k[7:]  # strip 'module.' if present
-                new_state[name] = v
+            for k, v in ckpt[key].items():
+                new_state[k[7:]] = v
             model[key].load_state_dict(new_state, strict=False)
-_ = [model[k].eval() for k in model]
 
-# Diffusion sampler
 sampler = DiffusionSampler(
     model.diffusion.diffusion,
     sampler=ADPM2Sampler(),
@@ -103,115 +80,211 @@ sampler = DiffusionSampler(
     clamp=False,
 )
 
-# -----------------------------
-# Audio helper: mel preprocessing
-# -----------------------------
-_to_mel = torchaudio.transforms.MelSpectrogram(
-    n_mels=80, n_fft=2048, win_length=1200, hop_length=300
+# ==============================================================
+# Phonemizer
+# ==============================================================
+import phonemizer
+vi_phonemizer = phonemizer.backend.EspeakBackend(
+    language="vi", preserve_punctuation=True, with_stress=True
 )
-_MEAN, _STD = -4.0, 4.0
+
+def phonemize_text(text: str) -> str:
+    ps = vi_phonemizer.phonemize([text])[0]
+    return ps.replace("(en)", "").replace("(vi)", "").strip()
 
 def length_to_mask(lengths: torch.LongTensor) -> torch.Tensor:
     mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
     mask = torch.gt(mask + 1, lengths.unsqueeze(1))
     return mask
 
-def preprocess(wave: np.ndarray) -> torch.Tensor:
-    """Same name as your snippet: np.float -> mel (normed)"""
-    wave_tensor = torch.from_numpy(wave).float()
-    mel_tensor = _to_mel(wave_tensor)
-    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - _MEAN) / _STD
-    return mel_tensor
-
-# -----------------------------
-# Phonemizer (vi)
-# -----------------------------
-import phonemizer
-vi_phonemizer = phonemizer.backend.EspeakBackend(language="vi", preserve_punctuation=True, with_stress=True)
-global_phonemizer = vi_phonemizer
-
-def phonemize_text(text: str) -> str:
-    ps = global_phonemizer.phonemize([text])[0]
-    return ps.replace("(en)", "").replace("(vi)", "").strip()
-
-# -----------------------------
-# Style extractor (from file path)
-# -----------------------------
-def compute_style(model, path, device):
-    """Compute style/prosody reference from a wav file path"""
-    wave, sr = librosa.load(path, sr=None, mono=True)
-    audio, _ = librosa.effects.trim(wave, top_db=30)
-    if sr != SR_OUT:
-        audio = librosa.resample(audio, sr, SR_OUT)
-    mel_tensor = preprocess(audio).to(device)
-
-    with torch.no_grad():
-        ref_s = model.style_encoder(mel_tensor.unsqueeze(1))
-        ref_p = model.predictor_encoder(mel_tensor.unsqueeze(1))
-    return torch.cat([ref_s, ref_p], dim=1)   # [1, 256]
-
-# Style extractor (from numpy array)
-def compute_style_from_numpy(model, arr: np.ndarray, sr: int, device):
-    if arr.ndim > 1:
-        arr = librosa.to_mono(arr.T)
-    audio, _ = librosa.effects.trim(arr, top_db=30)
-    if sr != SR_OUT:
-        audio = librosa.resample(audio, sr, SR_OUT)
-    mel_tensor = preprocess(audio).to(device)
-    with torch.no_grad():
-        ref_s = model.style_encoder(mel_tensor.unsqueeze(1))
-        ref_p = model.predictor_encoder(mel_tensor.unsqueeze(1))
-    return torch.cat([ref_s, ref_p], dim=1)
-
-# -----------------------------
-# Inference (your exact logic)
-# -----------------------------
-# Tunables (can expose to UI later)
-ALPHA = 0.3
-BETA  = 0.7
-DIFFUSION_STEPS = 5
-EMBEDDING_SCALE = 1.0
-
-def inference_one(text, ref_feat, ipa_text=None,
-                  alpha=ALPHA, beta=BETA, diffusion_steps=DIFFUSION_STEPS, embedding_scale=EMBEDDING_SCALE):
-    # text -> phonemes -> tokens
-    ps = ipa_text if ipa_text is not None else phonemize_text(text)
+# ==============================================================
+# Load style embeddings đã tính sẵn
+# ==============================================================
+STYLE_PTH = "Models/styles_speaker_parallel.pth"
+print(f"Loading precomputed styles: {STYLE_PTH}")
+styles_dict = torch.load(STYLE_PTH, map_location=DEVICE)
+
+# fallback speaker nếu mixer rỗng
+SPEAKER_ORDER_PREF = ["Kore", "Puck", "Algenib", "Leda"]
+DEFAULT_SPK = next((s for s in SPEAKER_ORDER_PREF if s in styles_dict), list(styles_dict.keys())[0])
+
+def get_style_by_length(speaker: str, phoneme_len: int):
+    spk_tensor = styles_dict[speaker]  # [510, 1, 256] hoặc [510, 256]
+    idx = min(max(phoneme_len, 1), spk_tensor.shape[0]) - 1
+    feat = spk_tensor[idx]
+    # ép về [1,256]
+    if feat.ndim == 3:  # [1,1,256]
+        feat = feat.squeeze(0)
+    if feat.ndim == 2:  # [1,256]
+        feat = feat.squeeze(0)
+    return feat.unsqueeze(0).to(DEVICE)  # [1,256]
+
+# ==============================================================
+# Style mixing utils
+# ==============================================================
+def parse_mix_spec(spec: str) -> dict:
+    """Parse 'Kore:0.75,Puck:0.25' -> {'Kore':0.75,'Puck':0.25} (lọc lỗi, gộp trùng)."""
+    mix = {}
+    if not spec or not isinstance(spec, str):
+        return mix
+    for part in spec.split(","):
+        if ":" not in part:
+            continue
+        k, v = part.split(":", 1)
+        k = (k or "").strip()
+        if not k:
+            continue
+        try:
+            w = float((v or "").strip())
+        except Exception:
+            continue
+        if not np.isfinite(w) or w <= 0:
+            continue
+        mix[k] = mix.get(k, 0.0) + w
+    return mix
+
+def get_style_mixed_by_length(mix_dict: dict, phoneme_len: int):
+    """Trộn style của nhiều speaker theo trọng số. Trả về [1,256] trên DEVICE."""
+    if not mix_dict:
+        return get_style_by_length(DEFAULT_SPK, phoneme_len)
+
+    total = sum(max(0.0, float(w)) for w in mix_dict.values())
+    if total <= 0:
+        return get_style_by_length(DEFAULT_SPK, phoneme_len)
+
+    mix_feat = None
+    for spk, w in mix_dict.items():
+        if spk not in styles_dict:
+            print(f"[WARN] Speaker '{spk}' không có trong styles_dict, bỏ qua.")
+            continue
+        feat_i = get_style_by_length(spk, phoneme_len)  # [1,256]
+        wi = float(w) / total
+        mix_feat = feat_i * wi if mix_feat is None else mix_feat + feat_i * wi
+
+    if mix_feat is None:
+        return get_style_by_length(DEFAULT_SPK, phoneme_len)
+    return mix_feat  # [1,256]
+
+# ==============================================================
+# Audio postprocess (librosa): trim + denoise + remove internal silence
+# ==============================================================
+def _simple_spectral_denoise(y, sr, n_fft=1024, hop=256, prop_decrease=0.8):
+    if y.size == 0:
+        return y
+    D = librosa.stft(y, n_fft=n_fft, hop_length=hop, win_length=n_fft)
+    S = np.abs(D)
+    noise = np.median(S, axis=1, keepdims=True)
+    S_clean = S - prop_decrease * noise
+    S_clean = np.maximum(S_clean, 0.0)
+    gain = S_clean / (S + 1e-8)
+    D_denoised = D * gain
+    y_out = librosa.istft(D_denoised, hop_length=hop, win_length=n_fft, length=len(y))
+    return y_out
+
+def _concat_with_crossfade(segments, crossfade_samples=0):
+    if not segments:
+        return np.array([], dtype=np.float32)
+    out = segments[0].astype(np.float32, copy=True)
+    for seg in segments[1:]:
+        seg = seg.astype(np.float32, copy=False)
+        if crossfade_samples > 0 and out.size > 0 and seg.size > 0:
+            cf = min(crossfade_samples, out.size, seg.size)
+            fade_out = np.linspace(1.0, 0.0, cf, dtype=np.float32)
+            fade_in  = 1.0 - fade_out
+            tail = out[-cf:] * fade_out + seg[:cf] * fade_in
+            out = np.concatenate([out[:-cf], tail, seg[cf:]], axis=0)
+        else:
+            out = np.concatenate([out, seg], axis=0)
+    return out
+
+def _reduce_internal_silence(y, sr, top_db=30, min_keep_ms=40, crossfade_ms=8):
+    if y.size == 0:
+        return y
+    intervals = librosa.effects.split(y, top_db=top_db)
+    if intervals.size == 0:
+        return y
+    min_keep = int(sr * (min_keep_ms / 1000.0))
+    segs = []
+    for s, e in intervals:
+        if e - s >= min_keep:
+            segs.append(y[s:e])
+    if not segs:
+        return y
+    crossfade = int(sr * (crossfade_ms / 1000.0))
+    y_out = _concat_with_crossfade(segs, crossfade_samples=crossfade)
+    return y_out
+
+def postprocess_audio(y, sr,
+                      trim_top_db=30,
+                      denoise=True,
+                      denoise_n_fft=1024,
+                      denoise_hop=256,
+                      denoise_strength=0.8,
+                      remove_internal_silence=True,
+                      split_top_db=30,
+                      min_keep_ms=40,
+                      crossfade_ms=8):
+    if y.size == 0:
+        return y.astype(np.float32)
+    y_trim, _ = librosa.effects.trim(y, top_db=trim_top_db)
+    if denoise:
+        y_trim = _simple_spectral_denoise(
+            y_trim, sr, n_fft=denoise_n_fft, hop=denoise_hop, prop_decrease=denoise_strength
+        )
+    if remove_internal_silence:
+        y_trim = _reduce_internal_silence(
+            y_trim, sr, top_db=split_top_db, min_keep_ms=min_keep_ms, crossfade_ms=crossfade_ms
+        )
+    y_trim = np.nan_to_num(y_trim, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
+    m = np.max(np.abs(y_trim)) + 1e-8
+    if m > 1.0:
+        y_trim = y_trim / m
+    return y_trim
+
+# ==============================================================
+# Inference core
+# ==============================================================
+def inference_one(text, ref_feat, alpha=ALPHA, beta=BETA,
+                  diffusion_steps=DIFFUSION_STEPS, embedding_scale=EMBEDDING_SCALE):
+    ps = phonemize_text(text)
     tokens = textcleaner(ps)
-    tokens.insert(0, 0)  # prepend BOS
-    tokens = torch.LongTensor(tokens).to(DEVICE).unsqueeze(0)  # [1, T]
+    tokens.insert(0, 0)
+    tokens = torch.LongTensor(tokens).unsqueeze(0).to(DEVICE)
+    input_lengths = torch.LongTensor([tokens.shape[-1]]).to(DEVICE)
+    text_mask = length_to_mask(input_lengths).to(DEVICE)
 
     with torch.no_grad():
-        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(DEVICE)
-        text_mask = length_to_mask(input_lengths).to(DEVICE)
-
-        # encoders
         t_en   = model.text_encoder(tokens, input_lengths, text_mask)
         bert_d = model.bert(tokens, attention_mask=(~text_mask).int())
         d_en   = model.bert_encoder(bert_d).transpose(-1, -2)
 
-        # diffusion for style latent
-        s_pred = sampler(
-            noise=torch.randn((1, 256)).unsqueeze(1).to(DEVICE),
-            embedding=bert_d,
-            embedding_scale=embedding_scale,
-            features=ref_feat,   # [1, 256]
-            num_steps=diffusion_steps,
-        ).squeeze(1)  # [1, 256]
-
-        s   = s_pred[:, 128:]    # prosody
-        ref = s_pred[:, :128]    # timbre
-
-        # blend with real ref features
+        if alpha == 0 and beta == 0:
+            s_pred = ref_feat.clone()  # [1,256]
+        else:
+            s_pred = sampler(
+                noise=torch.randn((1, 256)).unsqueeze(1).to(DEVICE),
+                embedding=bert_d,
+                embedding_scale=embedding_scale,
+                features=ref_feat,   # [1,256]
+                num_steps=diffusion_steps,
+            ).squeeze(1)  # [1,256]
+            
+        s, ref = s_pred[:, 128:], s_pred[:, :128]
         ref = alpha * ref + (1 - alpha) * ref_feat[:, :128]
         s   = beta  * s   + (1 - beta)  * ref_feat[:, 128:]
 
-        # duration prediction
+        # --- Metrics (cosine) ---
+        def cosine_sim(a, b):
+            return torch.nn.functional.cosine_similarity(a, b, dim=1).mean().item()
+        simi_timbre  = cosine_sim(s_pred[:, :128], ref_feat[:, :128])
+        simi_prosody = cosine_sim(s_pred[:, 128:], ref_feat[:, 128:])
+
+        # --- Duration / Alignment ---
         d = model.predictor.text_encoder(d_en, s, input_lengths, text_mask)
         x, _ = model.predictor.lstm(d)
         duration = torch.sigmoid(model.predictor.duration_proj(x)).sum(axis=-1)
         pred_dur = torch.round(duration.squeeze()).clamp(min=1)
 
-        # alignment
         T = int(pred_dur.sum().item())
         pred_aln = torch.zeros(input_lengths.item(), T, device=DEVICE)
         c = 0
@@ -220,119 +293,225 @@ def inference_one(text, ref_feat, ipa_text=None,
             pred_aln[i, c:c+span] = 1.0
             c += span
 
-        # prosody enc
         en = (d.transpose(-1, -2) @ pred_aln.unsqueeze(0))
         if model_params.decoder.type == "hifigan":
-            asr_new = torch.zeros_like(en); asr_new[:, :, 0] = en[:, :, 0]; asr_new[:, :, 1:] = en[:, :, 0:-1]; en = asr_new
+            en = torch.cat([en[:, :, :1], en[:, :, :-1]], dim=2)
 
         F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
-
-        # content (ASR-aligned)
         asr = (t_en @ pred_aln.unsqueeze(0))
         if model_params.decoder.type == "hifigan":
-            asr_new = torch.zeros_like(asr); asr_new[:, :, 0] = asr[:, :, 0]; asr_new[:, :, 1:] = asr[:, :, 0:-1]; asr = asr_new
+            asr = torch.cat([asr[:, :, :1], asr[:, :, :-1]], dim=2)
 
-        # decode
         out = model.decoder(asr, F0_pred, N_pred, ref.squeeze().unsqueeze(0))
 
     wav = out.squeeze().detach().cpu().numpy()
     if wav.shape[-1] > 50:
-        wav = wav[..., :-50]
-    return wav, ps
-
-# -----------------------------
+        wav = wav[:-50]
+
+    # Hậu xử lý: trim + denoise + bỏ silence nội bộ
+    wav = postprocess_audio(
+        wav, SR_OUT,
+        trim_top_db=30,
+        denoise=True,
+        denoise_n_fft=1024, denoise_hop=256, denoise_strength=0.8,
+        remove_internal_silence=True,
+        split_top_db=30, min_keep_ms=40, crossfade_ms=8
+    )
+    return wav, ps, simi_timbre, simi_prosody
+
+# ==============================================================
+# Ref-audio mapping (quét ./ref_voice để tìm file mẫu theo speaker)
+# ==============================================================
+def _norm(s: str) -> str:
+    import unicodedata
+    s = unicodedata.normalize("NFKD", s)
+    s = "".join([c for c in s if not unicodedata.combining(c)])
+    s = s.lower()
+    s = re.sub(r"[^a-z0-9_\-\.]+", "", s)
+    return s
+
+def build_ref_map(ref_dir: str) -> dict:
+    paths = glob.glob(os.path.join(ref_dir, "**", "*.wav"), recursive=True)
+    by_name = {}
+    for p in paths:
+        fname = os.path.basename(p)
+        by_name[_norm(fname)] = p
+
+    spk_map = {}
+    speakers = list(styles_dict.keys()) if isinstance(styles_dict, dict) else ["Kore","Algenib","Puck","Leda"]
+
+    for spk in speakers:
+        spk_n = _norm(spk)
+        hit = None
+        for k, p in by_name.items():
+            if f"_{spk_n}_" in k:
+                hit = p
+                break
+        if not hit:
+            for k, p in by_name.items():
+                if spk_n in k:
+                    hit = p
+                    break
+        if hit:
+            spk_map[spk] = hit
+    return spk_map
+
+REF_MAP = build_ref_map(REF_DIR)
+
+def get_ref_path_for_speaker(spk: str):
+    return REF_MAP.get(spk)
+
+# ==============================================================
+# Wrapper cho Gradio (nhận speaker_mix_spec là string ẩn)
+# ==============================================================
+def run_inference(text, alpha, beta, speaker_mix_spec):
+    normalizer = SoeNormalizer()
+    text = normalizer.normalize(text).replace(" ,", ",").replace(" .", ".")
+
+    ps = phonemize_text(text)
+    phoneme_len = len(ps.replace(" ", ""))
+
+    mix_dict = parse_mix_spec(speaker_mix_spec)
+    if len(mix_dict) > 0:
+        ref_feat = get_style_mixed_by_length(mix_dict, phoneme_len)
+        ref_idx = min(phoneme_len, 510)
+        total = sum(mix_dict.values())
+        mix_info = {k: round(float(v / total), 3) for k, v in mix_dict.items()}
+        chosen_speakers = list(mix_dict.keys())
+    else:
+        ref_feat = get_style_by_length(DEFAULT_SPK, phoneme_len)
+        ref_idx = min(phoneme_len, 510)
+        mix_info = {DEFAULT_SPK: 1.0}
+        chosen_speakers = [DEFAULT_SPK]
+
+    t0 = time.time()
+    wav, ps_out, simi_timbre, simi_prosody = inference_one(
+        text, ref_feat, alpha=float(alpha), beta=float(beta)
+    )
+    gen_time = time.time() - t0
+    rtf = gen_time / max(1e-6, len(wav) / SR_OUT)
+
+    info = {
+        "Text after soe_vinorms:": text,
+        "Speakers": chosen_speakers,
+        "Mix weights (normalized)": mix_info,
+        "Phonemes": ps_out,
+        "Phoneme length": phoneme_len,
+        "Ref index": ref_idx,
+        "simi_timbre": round(float(simi_timbre), 4),
+        "simi_prosody": round(float(simi_prosody), 4),
+        "alpha": float(alpha),
+        "beta": float(beta),
+        "RTF": round(float(rtf), 3),
+        "Device": DEVICE,
+    }
+    return (SR_OUT, wav.astype(np.float32)), info
+
+# ==============================================================
+# UI helper: build mix-spec CỐ ĐỊNH theo 4 speaker
+# ==============================================================
+def _build_mix_spec_ui_fixed(normalize, w1, w2, w3, w4, order):
+    pairs = [(order[0], float(w1 or 0.0)),
+             (order[1], float(w2 or 0.0)),
+             (order[2], float(w3 or 0.0)),
+             (order[3], float(w4 or 0.0))]
+    pairs = [(s, w) for s, w in pairs if w > 0]
+
+    if not pairs:
+        return "", {}, "**Sum:** 0.000"
+
+    total = sum(w for _, w in pairs)
+    if normalize and total > 0:
+        pairs = [(s, w/total) for s, w in pairs]
+
+    acc = {}
+    for s, w in pairs:
+        acc[s] = acc.get(s, 0.0) + w
+
+    mix_spec = ",".join([f"{s}:{w:.4f}" for s, w in acc.items()])
+    mix_view = {"weights": {s: round(w, 3) for s, w in acc.items()}, "normalized": bool(normalize)}
+    sum_md   = f"**Sum:** {round(sum(acc.values()), 3)}"
+    return mix_spec, mix_view, sum_md
+
+
+
+# ==============================================================
 # Gradio UI
-# -----------------------------
-
-
-SR_OUT = 24000
-ROOT_REF = "ref_voice"
-EXTS = {".wav", ".mp3", ".flac", ".ogg", ".m4a"}
-
-# -------- scan ref_voice/<id>_<speaker>/*.wav --------
-def scan_ref_voice(root=ROOT_REF):
-    """
-    return:
-      speakers: list[str]                # ví dụ: ["0_Fonos.vn", "1_James_A._Robinson", ...]
-      files_by_spk: dict[str, list[str]] # speaker_dir -> [full_path,...]
-    """
-    speakers, files_by_spk = [], {}
-    if not os.path.isdir(root):
-        return speakers, files_by_spk
-
-    for spk_dir in sorted(os.listdir(root)):
-        full_dir = os.path.join(root, spk_dir)
-        if not os.path.isdir(full_dir) or spk_dir.startswith("."):
-            continue
-        lst = []
-        for fn in sorted(os.listdir(full_dir)):
-            if os.path.splitext(fn)[1].lower() in EXTS:
-                lst.append(os.path.join(full_dir, fn))
-        if lst:
-            speakers.append(spk_dir)
-            files_by_spk[spk_dir] = lst
-    return speakers, files_by_spk
-
-SPEAKERS, FILES_BY_SPK = scan_ref_voice()
-
-with gr.Blocks(title="StyleTTS2-vi Demo ✨") as demo:
-    gr.Markdown("# StyleTTS2-vi Demo ✨")
+# ==============================================================
+with gr.Blocks(title="StyleTTS2-vi Demo") as demo:
+    gr.Markdown("# StyleTTS2-vi Demo")
 
     with gr.Row():
         with gr.Column():
-            text_inp = gr.Textbox(label="Text", lines=4,
-                                  value="Thời tiết hôm nay tại Hà Nội, nhiệt độ khoảng 27 độ C, có nắng nhẹ, rất hợp lý để mình đi dạo công viên nhé.")
-
-            # --- 1 ô audio duy nhất (nhận filepath) ---
-            ref_audio = gr.Audio(
-                label="Reference Audio",
-                type="filepath",                 # nhận đường dẫn file
-                sources=["upload","microphone"], # vẫn cho upload/mic
-                interactive=True,
-            )
-            ref_path  = gr.Textbox(label="Đường dẫn reference", interactive=False)
-
-            # --- chọn speaker -> hiện file tương ứng ---
-            spk_dd = gr.Dropdown(
-                label="Speaker",
-                choices=["(None)"] + SPEAKERS,
-                value="(None)",
-            )
-            file_dd = gr.Dropdown(
-                label="Voice in speaker",
-                choices=["(None)"],
-                value="(None)",
+            text_inp = gr.Textbox(
+                label="Text",
+                lines=4,
+                value="Trăng treo lơ lửng trên đỉnh núi chơ vơ, ánh sáng bàng bạc phủ lên bãi đá ngổn ngang. Con dế thổn thức trong khe cỏ, tiếng gió hun hút lùa qua hốc núi trập trùng. Dưới thung lũng, đàn trâu gặm cỏ ung dung, hơi sương vẩn đục, lảng bảng giữa đồng khuya tĩnh mịch."
             )
 
-            # khi chọn speaker -> cập nhật danh sách file
-            def on_pick_speaker(spk):
-                if spk == "(None)":
-                    return gr.update(choices=["(None)"], value="(None)")
-                files = FILES_BY_SPK.get(spk, [])
-                # hiển thị chỉ tên file cho gọn
-                labels = [os.path.basename(p) for p in files]
-                # ta sẽ map label->path bằng index; set value = mục đầu tiên
-                return gr.update(choices=labels, value=(labels[0] if labels else "(None)"))
-
-            spk_dd.change(on_pick_speaker, inputs=spk_dd, outputs=file_dd)
-
-            # map label (basename) -> full path theo speaker hiện tại
-            def on_pick_file(spk, label):
-                if spk == "(None)" or label == "(None)":
-                    return gr.update(value=None), ""
-                files = FILES_BY_SPK.get(spk, [])
-                # tìm đúng file theo basename
-                for p in files:
-                    if os.path.basename(p) == label:
-                        return gr.update(value=p), p  # set vào Audio + hiển thị path
-                return gr.update(value=None), ""
-
-            file_dd.change(on_pick_file, inputs=[spk_dd, file_dd], outputs=[ref_audio, ref_path])
-
-            # nếu người dùng upload/mic thì hiển thị luôn đường dẫn file tạm
-            def on_audio_changed(fp):
-                return fp or ""
-            ref_audio.change(on_audio_changed, inputs=ref_audio, outputs=ref_path)
+            # Danh sách speaker có trong styles_dict
+            spk_choices = list(styles_dict.keys()) if isinstance(styles_dict, dict) else ["Kore","Algenib","Puck","Leda"]
+
+            # Thứ tự CỐ ĐỊNH cho mixer
+            fixed_order = [s for s in ["Kore", "Puck", "Algenib", "Leda"] if s in spk_choices]
+            if len(fixed_order) < 4:
+                for s in spk_choices:
+                    if s not in fixed_order:
+                        fixed_order.append(s)
+                    if len(fixed_order) == 4:
+                        break
+
+            # === Luôn hiển thị 4 voice sample ===
+            with gr.Accordion("Reference samples", open=True):
+                with gr.Row():
+                    spk0 = fixed_order[0] if len(fixed_order) > 0 else "Kore"
+                    spk1 = fixed_order[1] if len(fixed_order) > 1 else "Puck"
+                    with gr.Column():
+                        gr.Markdown(f"**{spk0}**")
+                        gr.Audio(value=get_ref_path_for_speaker(spk0), label=f"{spk0} sample", type="filepath", interactive=False)
+                    with gr.Column():
+                        gr.Markdown(f"**{spk1}**")
+                        gr.Audio(value=get_ref_path_for_speaker(spk1), label=f"{spk1} sample", type="filepath", interactive=False)
+                with gr.Row():
+                    spk2 = fixed_order[2] if len(fixed_order) > 2 else "Algenib"
+                    spk3 = fixed_order[3] if len(fixed_order) > 3 else "Leda"
+                    with gr.Column():
+                        gr.Markdown(f"**{spk2}**")
+                        gr.Audio(value=get_ref_path_for_speaker(spk2), label=f"{spk2} sample", type="filepath", interactive=False)
+                    with gr.Column():
+                        gr.Markdown(f"**{spk3}**")
+                        gr.Audio(value=get_ref_path_for_speaker(spk3), label=f"{spk3} sample", type="filepath", interactive=False)
+
+            # ---- Style Mixer cố định 4 slot ----
+            with gr.Accordion("Style Mixer", open=True):
+                normalize_ck = gr.Checkbox(value=True, label="Normalize weights to 1")
+
+                # Hàng 1: Kore & Puck
+                with gr.Row(equal_height=True):
+                    with gr.Column():
+                        gr.Markdown(f"**{fixed_order[0]}**")
+                        w1 = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Weight 1", container=False)
+                    with gr.Column():
+                        gr.Markdown(f"**{fixed_order[1]}**")
+                        w2 = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Weight 2", container=False)
+
+                # Hàng 2: Algenib & Leda
+                with gr.Row(equal_height=True):
+                    with gr.Column():
+                        gr.Markdown(f"**{fixed_order[2]}**")
+                        w3 = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Weight 3", container=False)
+                    with gr.Column():
+                        gr.Markdown(f"**{fixed_order[3]}**")
+                        w4 = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Weight 4", container=False)
+
+                mix_sum_md     = gr.Markdown("**Sum:** 0.000")
+                mix_view_json  = gr.JSON(label="Mixer weights (view)")
+                mix_spec_state = gr.State("")          # string mix-spec cho backend
+                order_state    = gr.State(fixed_order) # giữ thứ tự cố định
+
+            with gr.Row():
+                alpha_n = gr.Number(value=ALPHA, label="alpha diffusion (0-1, timbre)", precision=3)
+                beta_n  = gr.Number(value=BETA,  label="beta diffusion (0-1, prosody)", precision=3)
 
             btn = gr.Button("Đọc 🔊🔥", variant="primary")
 
@@ -340,35 +519,25 @@ with gr.Blocks(title="StyleTTS2-vi Demo ✨") as demo:
             out_audio = gr.Audio(label="Synthesised Audio", type="numpy")
             metrics   = gr.JSON(label="Metrics")
 
-    # ---- Inference: xử lý từ filepath ----
-    def _run(text, ref_fp):
-        # ref_fp là string path (do type='filepath')
-        if isinstance(ref_fp, str) and os.path.isfile(ref_fp):
-            wav, _ = librosa.load(ref_fp, sr=SR_OUT, mono=True)
-            ref_feat = compute_style_from_numpy(model, wav, SR_OUT, DEVICE)
-            ref_src = ref_fp
-        else:
-            ref_feat = torch.zeros(1, 256).to(DEVICE)
-            ref_src = "(None)"
-
-        t0 = time.time()
-        wav, ps = inference_one(text, ref_feat)
-        wav = wav.astype(np.float32)
-        gen_time = time.time() - t0
-        rtf = gen_time / max(1e-6, len(wav)/SR_OUT)
-
-        info = {
-            "Ref path": ref_src,
-            "Phonemes": ps,
-            "Sample rate": SR_OUT,
-            "RTF": round(float(rtf), 3),
-            "Device": DEVICE,
-        }
-        return (SR_OUT, wav), info
-
-    btn.click(_run, inputs=[text_inp, ref_audio], outputs=[out_audio, metrics])
-
-
+    # Bất kỳ thay đổi weight/normalize -> build spec cố định + update tổng/json
+    def _ui_build_wrapper_fixed(normalize, w1, w2, w3, w4, order):
+        spec, view, summ = _build_mix_spec_ui_fixed(normalize, w1, w2, w3, w4, order)
+        return spec, view, summ
+
+    for comp in [normalize_ck, w1, w2, w3, w4]:
+        comp.change(
+            _ui_build_wrapper_fixed,
+            inputs=[normalize_ck, w1, w2, w3, w4, order_state],
+            outputs=[mix_spec_state, mix_view_json, mix_sum_md]
+        )
+
+
+    # Nút đọc: dùng mix_spec_state; nếu rỗng => fallback DEFAULT_SPK
+    btn.click(
+        run_inference,
+        inputs=[text_inp, alpha_n, beta_n, mix_spec_state],
+        outputs=[out_audio, metrics]
+    )
 
 if __name__ == "__main__":
     demo.launch()

app2.py

@@ -0,0 +1,419 @@
+# -*- coding: utf-8 -*-
+"""
+Gradio app.py - wired to your 'inference_one' implementation
+- Reference voice: upload OR choose from train_ref/
+- Uses phonemize_text(), compute_style(), inference_one() exactly like your snippet
+- NOW: adds UI sliders for alpha and beta and threads them into inference
+"""
+
+import os
+import time
+import yaml
+import numpy as np
+import torch
+import torchaudio
+import librosa
+import gradio as gr
+
+from munch import Munch
+
+# -----------------------------
+# Reproducibility
+# -----------------------------
+torch.manual_seed(0)
+torch.backends.cudnn.benchmark = False
+torch.backends.cudnn.deterministic = True
+np.random.seed(0)
+
+# -----------------------------
+# Device / sample-rate
+# -----------------------------
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+SR_OUT = 24000  # target audio rate for synthesis
+
+# -----------------------------
+# External modules from the project
+# -----------------------------
+from models import *             # noqa: F401,F403
+from utils import *              # noqa: F401,F403
+from models import build_model
+from text_utils import TextCleaner
+from Utils_extend_v1.PLBERT.util import load_plbert
+from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule
+
+textcleaner = TextCleaner()
+
+# -----------------------------
+# Config / model loading
+# -----------------------------
+from huggingface_hub import hf_hub_download
+hf_hub_download(
+    repo_id="ltphuongunited/styletts2_vi", 
+    filename="epoch_2nd_00058.pth", 
+    local_dir="Models/multi_phoaudio_gemini", 
+    local_dir_use_symlinks=False,
+)
+
+# CONFIG_PATH = os.getenv("MODEL_CONFIG", "Models/multi_phoaudio_gemini/config_phoaudio_gemini_small.yml")
+# CHECKPOINT_PTH = os.getenv("MODEL_CKPT", "Models/multi_phoaudio_gemini/epoch_2nd_00058.pth")
+
+CHECKPOINT_PTH = "Models/gemini_vi/gemini_2nd_00045.pth"
+CONFIG_PATH = "Models/gemini_vi/config_gemini_vi_en.yml"
+
+# Load config
+config = yaml.safe_load(open(CONFIG_PATH))
+
+# Build components
+ASR_config = config.get("ASR_config", False)
+ASR_path   = config.get("ASR_path", False)
+F0_path    = config.get("F0_path", False)
+PLBERT_dir = config.get("PLBERT_dir", False)
+
+text_aligner    = load_ASR_models(ASR_path, ASR_config)
+pitch_extractor = load_F0_models(F0_path)
+plbert          = load_plbert(PLBERT_dir)
+
+model_params = recursive_munch(config["model_params"])
+model = build_model(model_params, text_aligner, pitch_extractor, plbert)
+
+# to device & eval
+_ = [model[k].to(DEVICE) for k in model]
+_ = [model[k].eval() for k in model]
+
+# Load checkpoint
+if not os.path.isfile(CHECKPOINT_PTH):
+    raise FileNotFoundError(f"Checkpoint not found at '{CHECKPOINT_PTH}'")
+ckpt = torch.load(CHECKPOINT_PTH, map_location="cpu")
+params = ckpt["net"]
+for key in model:
+    if key in params:
+        try:
+            model[key].load_state_dict(params[key])
+        except Exception:
+            from collections import OrderedDict
+            state_dict = params[key]
+            new_state = OrderedDict()
+            for k, v in state_dict.items():
+                name = k[7:]  # strip 'module.' if present
+                new_state[name] = v
+            model[key].load_state_dict(new_state, strict=False)
+_ = [model[k].eval() for k in model]
+
+# Diffusion sampler
+sampler = DiffusionSampler(
+    model.diffusion.diffusion,
+    sampler=ADPM2Sampler(),
+    sigma_schedule=KarrasSchedule(sigma_min=1e-4, sigma_max=3.0, rho=9.0),
+    clamp=False,
+)
+
+# -----------------------------
+# Audio helper: mel preprocessing
+# -----------------------------
+_to_mel = torchaudio.transforms.MelSpectrogram(
+    n_mels=80, n_fft=2048, win_length=1200, hop_length=300
+)
+_MEAN, _STD = -4.0, 4.0
+
+def length_to_mask(lengths: torch.LongTensor) -> torch.Tensor:
+    mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
+    mask = torch.gt(mask + 1, lengths.unsqueeze(1))
+    return mask
+
+def preprocess(wave: np.ndarray) -> torch.Tensor:
+    """Same name as your snippet: np.float -> mel (normed)"""
+    wave_tensor = torch.from_numpy(wave).float()
+    mel_tensor = _to_mel(wave_tensor)
+    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - _MEAN) / _STD
+    return mel_tensor
+
+# -----------------------------
+# Phonemizer (vi)
+# -----------------------------
+import phonemizer
+vi_phonemizer = phonemizer.backend.EspeakBackend(language="vi", preserve_punctuation=True, with_stress=True)
+global_phonemizer = vi_phonemizer
+
+def phonemize_text(text: str) -> str:
+    ps = global_phonemizer.phonemize([text])[0]
+    return ps.replace("(en)", "").replace("(vi)", "").strip()
+
+# -----------------------------
+# Style extractor (from file path)
+# -----------------------------
+
+def compute_style(model, path, device):
+    """Compute style/prosody reference from a wav file path"""
+    wave, sr = librosa.load(path, sr=None, mono=True)
+    audio, _ = librosa.effects.trim(wave, top_db=30)
+    if sr != SR_OUT:
+        audio = librosa.resample(audio, sr, SR_OUT)
+    mel_tensor = preprocess(audio).to(device)
+
+    with torch.no_grad():
+        ref_s = model.style_encoder(mel_tensor.unsqueeze(1))
+        ref_p = model.predictor_encoder(mel_tensor.unsqueeze(1))
+    return torch.cat([ref_s, ref_p], dim=1)   # [1, 256]
+
+# Style extractor (from numpy array)
+
+def compute_style_from_numpy(model, arr: np.ndarray, sr: int, device):
+    if arr.ndim > 1:
+        arr = librosa.to_mono(arr.T)
+    audio, _ = librosa.effects.trim(arr, top_db=30)
+    if sr != SR_OUT:
+        audio = librosa.resample(audio, sr, SR_OUT)
+    mel_tensor = preprocess(audio).to(device)
+    with torch.no_grad():
+        ref_s = model.style_encoder(mel_tensor.unsqueeze(1))
+        ref_p = model.predictor_encoder(mel_tensor.unsqueeze(1))
+    return torch.cat([ref_s, ref_p], dim=1)
+
+# -----------------------------
+# Inference (your exact logic)
+# -----------------------------
+# Tunables (still as defaults; UI will override)
+ALPHA = 0.3
+BETA  = 0.7
+DIFFUSION_STEPS = 5
+EMBEDDING_SCALE = 1.0
+
+def inference_one(text, ref_feat, ipa_text=None,
+                  alpha=ALPHA, beta=BETA, diffusion_steps=DIFFUSION_STEPS, embedding_scale=EMBEDDING_SCALE):
+    # text -> phonemes -> tokens
+    ps = ipa_text if ipa_text is not None else phonemize_text(text)
+    tokens = textcleaner(ps)
+    tokens.insert(0, 0)  # prepend BOS
+    tokens = torch.LongTensor(tokens).to(DEVICE).unsqueeze(0)  # [1, T]
+
+    with torch.no_grad():
+        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(DEVICE)
+        text_mask = length_to_mask(input_lengths).to(DEVICE)
+
+        # encoders
+        t_en   = model.text_encoder(tokens, input_lengths, text_mask)
+        bert_d = model.bert(tokens, attention_mask=(~text_mask).int())
+        d_en   = model.bert_encoder(bert_d).transpose(-1, -2)
+
+
+        if alpha == 0 and beta == 0:
+            print("Ignore Diffusion")
+            ref = ref_feat[:, :128]
+            s = ref_feat[:, 128:]
+            simi_timbre, simi_prosody = 1,1
+        else:
+            print("Have Diffusion")
+            # diffusion for style latent
+            s_pred = sampler(
+                noise=torch.randn((1, 256)).unsqueeze(1).to(DEVICE),
+                embedding=bert_d,
+                embedding_scale=embedding_scale,
+                features=ref_feat,   # [1, 256]
+                num_steps=diffusion_steps,
+            ).squeeze(1)  # [1, 256]
+
+            s   = s_pred[:, 128:]    # prosody
+            ref = s_pred[:, :128]    # timbre
+
+            # blend with real ref features
+            ref = alpha * ref + (1 - alpha) * ref_feat[:, :128]
+            s   = beta  * s   + (1 - beta)  * ref_feat[:, 128:]
+
+            with torch.no_grad():
+                ref0 = ref_feat[:, :128]   # timbre gốc
+                s0   = ref_feat[:, 128:]   # prosody gốc
+
+                eps = 1e-8
+
+                def stats(name, new, base):
+                    delta = new - base
+                    l2_delta   = torch.norm(delta, dim=1)                          # ||Δ||
+                    l2_base    = torch.norm(base, dim=1) + eps                     # ||x||
+                    rel_l2     = (l2_delta / l2_base)                              # ||Δ|| / ||x||
+                    mae        = torch.mean(torch.abs(delta), dim=1)               # MAE
+                    cos_sim    = F.cosine_similarity(new, base, dim=1)             # cos(new, base)
+                    snr_db     = 20.0 * torch.log10(l2_base / (l2_delta + eps))    # SNR ~ 20*log10(||x||/||Δ||)
+                    # # Inference batch thường =1, nhưng vẫn in theo batch để tổng quát
+                    # for i in range(new.shape[0]):
+                    #     print(f"[{name}][sample {i}] "
+                    #         f"L2Δ={l2_delta[i]:.4f} | relL2={rel_l2[i]:.4f} | MAE={mae[i]:.6f} | "
+                    #         f"cos={cos_sim[i]:.4f} | SNR={snr_db[i]:.2f} dB")
+
+                    return cos_sim
+
+
+                simi_timbre = stats("REF(timbre)", s_pred[:, :128], ref_feat[:, :128]).detach().cpu().squeeze().item()
+                simi_prosody   = stats("S(prosody)",  s_pred[:, 128:],  ref_feat[:, 128:]).detach().cpu().squeeze().item()
+
+        # duration prediction
+        d = model.predictor.text_encoder(d_en, s, input_lengths, text_mask)
+        x, _ = model.predictor.lstm(d)
+        duration = torch.sigmoid(model.predictor.duration_proj(x)).sum(axis=-1)
+        pred_dur = torch.round(duration.squeeze()).clamp(min=1)
+
+        # alignment
+        T = int(pred_dur.sum().item())
+        pred_aln = torch.zeros(input_lengths.item(), T, device=DEVICE)
+        c = 0
+        for i in range(input_lengths.item()):
+            span = int(pred_dur[i].item())
+            pred_aln[i, c:c+span] = 1.0
+            c += span
+
+        # prosody enc
+        en = (d.transpose(-1, -2) @ pred_aln.unsqueeze(0))
+        if model_params.decoder.type == "hifigan":
+            asr_new = torch.zeros_like(en); asr_new[:, :, 0] = en[:, :, 0]; asr_new[:, :, 1:] = en[:, :, 0:-1]; en = asr_new
+
+        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
+
+        # content (ASR-aligned)
+        asr = (t_en @ pred_aln.unsqueeze(0))
+        if model_params.decoder.type == "hifigan":
+            asr_new = torch.zeros_like(asr); asr_new[:, :, 0] = asr[:, :, 0]; asr_new[:, :, 1:] = asr[:, :, 0:-1]; asr = asr_new
+
+        # decode
+        out = model.decoder(asr, F0_pred, N_pred, ref.squeeze().unsqueeze(0))
+
+    wav = out.squeeze().detach().cpu().numpy()
+    if wav.shape[-1] > 50:
+        wav = wav[..., :-50]
+    return wav, ps, simi_timbre, simi_prosody
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+
+SR_OUT = 24000
+ROOT_REF = "ref_voice"
+EXTS = {".wav", ".mp3", ".flac", ".ogg", ".m4a"}
+
+# -------- scan ref_voice/<id>_<speaker>/*.wav --------
+
+def scan_ref_voice(root=ROOT_REF):
+    """
+    return:
+      speakers: list[str]                # ví dụ: ["0_Fonos.vn", "1_James_A._Robinson", ...]
+      files_by_spk: dict[str, list[str]] # speaker_dir -> [full_path,...]
+    """
+    speakers, files_by_spk = [], {}
+    if not os.path.isdir(root):
+        return speakers, files_by_spk
+
+    for spk_dir in sorted(os.listdir(root)):
+        full_dir = os.path.join(root, spk_dir)
+        if not os.path.isdir(full_dir) or spk_dir.startswith("."):
+            continue
+        lst = []
+        for fn in sorted(os.listdir(full_dir)):
+            if os.path.splitext(fn)[1].lower() in EXTS:
+                lst.append(os.path.join(full_dir, fn))
+        if lst:
+            speakers.append(spk_dir)
+            files_by_spk[spk_dir] = lst
+    return speakers, files_by_spk
+
+SPEAKERS, FILES_BY_SPK = scan_ref_voice()
+
+with gr.Blocks(title="StyleTTS2-vi Demo ✨") as demo:
+    gr.Markdown("# StyleTTS2-vi Demo ✨")
+
+    with gr.Row():
+        with gr.Column():
+            text_inp = gr.Textbox(label="Text", lines=4,
+                                  value="Thời tiết hôm nay tại Hà Nội, nhiệt độ khoảng 27 độ C, có nắng nhẹ, rất hợp lý để mình đi dạo công viên nhé.")
+
+            # --- 1 ô audio duy nhất (nhận filepath) ---
+            ref_audio = gr.Audio(
+                label="Reference Audio",
+                type="filepath",                 # nhận đường dẫn file
+                sources=["upload","microphone"], # vẫn cho upload/mic
+                interactive=True,
+            )
+            ref_path  = gr.Textbox(label="Đường dẫn reference", interactive=False)
+
+            # --- chọn speaker -> hiện file tương ứng ---
+            spk_dd = gr.Dropdown(
+                label="Speaker",
+                choices=["(None)"] + SPEAKERS,
+                value="(None)",
+            )
+            file_dd = gr.Dropdown(
+                label="Voice in speaker",
+                choices=["(None)"],
+                value="(None)",
+            )
+
+            # khi chọn speaker -> cập nhật danh sách file
+            def on_pick_speaker(spk):
+                if spk == "(None)":
+                    return gr.update(choices=["(None)"], value="(None)")
+                files = FILES_BY_SPK.get(spk, [])
+                # hiển thị chỉ tên file cho gọn
+                labels = [os.path.basename(p) for p in files]
+                # ta sẽ map label->path bằng index; set value = mục đầu tiên
+                return gr.update(choices=labels, value=(labels[0] if labels else "(None)"))
+
+            spk_dd.change(on_pick_speaker, inputs=spk_dd, outputs=file_dd)
+
+            # map label (basename) -> full path theo speaker hiện tại
+            def on_pick_file(spk, label):
+                if spk == "(None)" or label == "(None)":
+                    return gr.update(value=None), ""
+                files = FILES_BY_SPK.get(spk, [])
+                # tìm đúng file theo basename
+                for p in files:
+                    if os.path.basename(p) == label:
+                        return gr.update(value=p), p  # set vào Audio + hiển thị path
+                return gr.update(value=None), ""
+
+            file_dd.change(on_pick_file, inputs=[spk_dd, file_dd], outputs=[ref_audio, ref_path])
+
+            # nếu người dùng upload/mic thì hiển thị luôn đường dẫn file tạm
+            def on_audio_changed(fp):
+                return fp or ""
+            ref_audio.change(on_audio_changed, inputs=ref_audio, outputs=ref_path)
+
+            # --- NEW: alpha/beta numeric inputs ---
+            with gr.Row():
+                alpha_n = gr.Number(value=ALPHA, label="alpha (0-1, timbre)", precision=3)
+                beta_n  = gr.Number(value=BETA,  label="beta (0-1, prosody)", precision=3)
+
+            btn = gr.Button("Đọc 🔊🔥", variant="primary")
+
+        with gr.Column():
+            out_audio = gr.Audio(label="Synthesised Audio", type="numpy")
+            metrics   = gr.JSON(label="Metrics")
+
+    # ---- Inference: xử lý từ filepath ----
+    def _run(text, ref_fp, alpha, beta):
+        # ref_fp là string path (do type='filepath')
+        if isinstance(ref_fp, str) and os.path.isfile(ref_fp):
+            wav, _ = librosa.load(ref_fp, sr=SR_OUT, mono=True)
+            ref_feat = compute_style_from_numpy(model, wav, SR_OUT, DEVICE)
+            ref_src = ref_fp
+        else:
+            ref_feat = torch.zeros(1, 256).to(DEVICE)
+            ref_src = "(None)"
+
+        t0 = time.time()
+        wav, ps, simi_timbre, simi_prosody = inference_one(text, ref_feat, alpha=float(alpha), beta=float(beta))
+        wav = wav.astype(np.float32)
+        gen_time = time.time() - t0
+        rtf = gen_time / max(1e-6, len(wav)/SR_OUT)
+
+        info = {
+            "simi_timbre": round(float(simi_timbre), 4) ,
+            "simi_prosody": round(float(simi_prosody), 4) ,
+            "Phonemes": ps,
+            "Sample rate": SR_OUT,
+            "RTF": round(float(rtf), 3),
+            "Device": DEVICE,
+        }
+        return (SR_OUT, wav), info
+
+    btn.click(_run, inputs=[text_inp, ref_audio, alpha_n, beta_n], outputs=[out_audio, metrics])
+
+
+if __name__ == "__main__":
+    demo.launch()

ref_voice/0000000_Kore_Quân_sự.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abfd3f8771395bbcb0789f2f5e61fab93ac186672d0ff756fbe5c44854bb4cc3
+size 284730

ref_voice/0000001_Algenib_Giáo_dục.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ddceacd307b7f1c0e9ba086bf3744a42c4fd128f94f1e50af9465c0d7a72eff
+size 411450

ref_voice/0000002_Puck_Giáo_dục.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec31b58f8ea3606eeb7c785d9bf66db952e15ad5fe8c74875771f29a37976dff
+size 760890

ref_voice/0000003_Leda_Giáo_dục.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fef8463fceb431353e8a6155cb571fe10827983cf72a85772822c6656bd72ba
+size 672570

requirements.txt

@@ -14,4 +14,5 @@ tqdm
 typing
 typing-extensions
 git+https://github.com/resemble-ai/monotonic_align.git
-phonemizer
+phonemizer
+soe-vinorm

train_second.py

@@ -349,7 +349,7 @@ def main(config_path):
             s_trg = torch.cat([gs, s_dur], dim=-1).detach() # ground truth for denoiser
 
             bert_dur = model.bert(texts, attention_mask=(~text_mask).int())
-            d_en = model.bert_encoder(bert_dur).transpose(-1, -2) 
+            d_en = model.bert_encodattention_masker(bert_dur).transpose(-1, -2) 
             
             # denoiser training
             if epoch >= diff_epoch: