src/f5_tts/model/utils.py

from __future__ import annotations

import os
import random
from collections import defaultdict
from importlib.resources import files
from typing import List, Union

import jieba
import torch
from pypinyin import Style, lazy_pinyin
from torch.nn.utils.rnn import pad_sequence
import requests
import json
import socket
import subprocess
from datetime import datetime


# seed everything


def seed_everything(seed=0):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


# helpers


def exists(v):
    return v is not None


def default(v, d):
    return v if exists(v) else d


def is_package_available(package_name: str) -> bool:
    try:
        import importlib

        package_exists = importlib.util.find_spec(package_name) is not None
        return package_exists
    except Exception:
        return False


# tensor helpers


def lens_to_mask(t: int["b"], length: int | None = None) -> bool["b n"]:  # noqa: F722 F821
    if not exists(length):
        length = t.amax()

    seq = torch.arange(length, device=t.device)
    return seq[None, :] < t[:, None]


def mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"]):  # noqa: F722 F821
    max_seq_len = seq_len.max().item()
    seq = torch.arange(max_seq_len, device=start.device).long()
    start_mask = seq[None, :] >= start[:, None]
    end_mask = seq[None, :] < end[:, None]
    return start_mask & end_mask


def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"]):  # noqa: F722 F821
    lengths = (frac_lengths * seq_len).long()
    max_start = seq_len - lengths

    rand = torch.rand_like(frac_lengths)
    start = (max_start * rand).long().clamp(min=0)
    end = start + lengths

    return mask_from_start_end_indices(seq_len, start, end)


def maybe_masked_mean(t: float["b n d"], mask: bool["b n"] = None) -> float["b d"]:  # noqa: F722
    if not exists(mask):
        return t.mean(dim=1)

    t = torch.where(mask[:, :, None], t, torch.tensor(0.0, device=t.device))
    num = t.sum(dim=1)
    den = mask.float().sum(dim=1)

    return num / den.clamp(min=1.0)


# simple utf-8 tokenizer, since paper went character based
def list_str_to_tensor(text: list[str], padding_value=-1) -> int["b nt"]:  # noqa: F722
    list_tensors = [torch.tensor([*bytes(t, "UTF-8")]) for t in text]  # ByT5 style
    text = pad_sequence(list_tensors, padding_value=padding_value, batch_first=True)
    return text


# char tokenizer, based on custom dataset's extracted .txt file
def list_str_to_idx(
    text: list[str] | list[list[str]],
    vocab_char_map: dict[str, int],  # {char: idx}
    padding_value=-1,
) -> int["b nt"]:  # noqa: F722
    list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text]  # pinyin or char style
    text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True)
    return text

# Get tokenizer

def _get_tokenizer(dataset_name, tokenizer, extra_vocab_path: str = None):
    """
    tokenizer   - "pinyin" do g2p for only chinese characters, need .txt vocab_file
                - "char" for char-wise tokenizer, need .txt vocab_file
                - "byte" for utf-8 tokenizer
                - "custom" if you're directly passing in a path to the vocab.txt you want to use
    vocab_size  - if use "pinyin", all available pinyin types, common alphabets (also those with accent) and symbols
                - if use "char", derived from unfiltered character & symbol counts of custom dataset
                - if use "byte", set to 256 (unicode byte range)
    extra_vocab_path - path to txt file containing additional characters (one per line) to expand vocabulary
    """
    
    if tokenizer in ["pinyin", "char", "cls"]:
        tokenizer_path = os.path.join(files("f5_tts").joinpath("../../data"), f"{dataset_name}_{tokenizer}/vocab.txt")
        with open(tokenizer_path, "r", encoding="utf-8") as f:
            vocab_char_map = {}
            for i, char in enumerate(f):
                vocab_char_map[char[:-1]] = i
        vocab_size = len(vocab_char_map)
        assert vocab_char_map[" "] == 0, "make sure space is of idx 0 in vocab.txt, cuz 0 is used for unknown char"
        
    elif tokenizer == "byte":
        vocab_char_map = None
        vocab_size = 256
        
    elif tokenizer == "custom":
        with open(dataset_name, "r", encoding="utf-8") as f:
            vocab_char_map = {}
            for i, char in enumerate(f):
                vocab_char_map[char[:-1]] = i
        vocab_size = len(vocab_char_map)
    
    # Load and merge extra vocabulary from txt file
    if extra_vocab_path is not None and os.path.exists(extra_vocab_path):
        if vocab_char_map is not None:  # Only extend if not byte tokenizer
            current_vocab_size = len(vocab_char_map)
            with open(extra_vocab_path, "r", encoding="utf-8") as f:
                for char in f:
                    char = char.strip()
                    if char and char not in vocab_char_map:
                        vocab_char_map[char] = len(vocab_char_map)
            vocab_size = len(vocab_char_map)
            print(f"Extended vocabulary with {vocab_size - current_vocab_size} new tokens from {extra_vocab_path}")
    
    return vocab_char_map, vocab_size



def get_tokenizer(dataset_name, tokenizer: str | List[str], extra_vocab_path: str = None):
    """
    Return a dictionary of tokenizers if tokenizer is a list, each key is the tokenizer name and value is a tuple of (vocab_char_map, vocab_size)
    Otherwise, return a dictionary with single tokenizer entry
    """
    
    if isinstance(tokenizer, list):
        tokenizers_dict = {}
        for t_name in tokenizer:
            vocab_char_map, vocab_size = _get_tokenizer(dataset_name, t_name, extra_vocab_path)
            tokenizers_dict[t_name] = (vocab_char_map, vocab_size)  # Fixed: was using 't' instead of 't_name'
        return tokenizers_dict
    else:
        vocab_char_map, vocab_size = _get_tokenizer(dataset_name, tokenizer, extra_vocab_path)
        return  vocab_char_map, vocab_size

def save_vocab(vocab_char_map, save_path):
    """Save vocabulary to file"""
    if vocab_char_map is None:
        return
        
    print(f"\nSaving vocabulary to: {save_path}")
    
    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    # Sort by index to maintain order
    vocab_items = sorted(vocab_char_map.items(), key=lambda x: x[1])
    
    with open(save_path, 'w', encoding='utf-8') as f:
        for char, idx in vocab_items:
            f.write(f"{char}\n")
    
    print(f"✓ Saved {len(vocab_char_map)} tokens to vocab.txt")

def send_slack_notification(message, webhook_url=None, title="Training Notification"):
    """Send a notification to a Slack channel via webhook."""
    if webhook_url is None:
        webhook_url = os.getenv("SLACK_WEBHOOK_URL")
    
    if not webhook_url:
        return False
    
    hostname = socket.gethostname()
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    payload = {
        "text": f"*{title}*\n*Time:* {timestamp}\n*Host:* {hostname}\n*Message:* {message}"
    }
    
    try:
        response = requests.post(
            webhook_url, 
            data=json.dumps(payload),
            headers={'Content-Type': 'application/json'}
        )
        return response.status_code == 200
    except Exception as e:
        print(f"Failed to send Slack notification: {e}")
        return False

def track_with_dvc(path):
    """Track a file or directory with DVC."""
    try:
        # Check if dvc is initialized
        if not os.path.exists(".dvc"):
            print("DVC not initialized. Skipping tracking.")
            return False
            
        print(f"Tracking with DVC: {path}")
        subprocess.run(["dvc", "add", path], check=True, capture_output=True)
        return True
    except subprocess.CalledProcessError as e:
        print(f"DVC add failed for {path}: {e.stderr.decode() if e.stderr else str(e)}")
        return False
    except Exception as e:
        print(f"Failed to track with DVC: {e}")
        return False

# convert char to pinyin


def convert_char_to_pinyin(text_list, polyphone=True):
    if jieba.dt.initialized is False:
        jieba.default_logger.setLevel(50)  # CRITICAL
        jieba.initialize()

    final_text_list = []
    custom_trans = str.maketrans(
        {";": ",", "“": '"', "”": '"', "‘": "'", "’": "'"}
    )  # add custom trans here, to address oov

    def is_chinese(c):
        return (
            "\u3100" <= c <= "\u9fff"  # common chinese characters
        )

    for text in text_list:
        char_list = []
        text = text.translate(custom_trans)
        for seg in jieba.cut(text):
            seg_byte_len = len(bytes(seg, "UTF-8"))
            if seg_byte_len == len(seg):  # if pure alphabets and symbols
                if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"":
                    char_list.append(" ")
                char_list.extend(seg)
            elif polyphone and seg_byte_len == 3 * len(seg):  # if pure east asian characters
                seg_ = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)
                for i, c in enumerate(seg):
                    if is_chinese(c):
                        char_list.append(" ")
                    char_list.append(seg_[i])
            else:  # if mixed characters, alphabets and symbols
                for c in seg:
                    if ord(c) < 256:
                        char_list.extend(c)
                    elif is_chinese(c):
                        char_list.append(" ")
                        char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))
                    else:
                        char_list.append(c)
        final_text_list.append(char_list)

    return final_text_list


# filter func for dirty data with many repetitions


def repetition_found(text, length=2, tolerance=10):
    pattern_count = defaultdict(int)
    for i in range(len(text) - length + 1):
        pattern = text[i : i + length]
        pattern_count[pattern] += 1
    for pattern, count in pattern_count.items():
        if count > tolerance:
            return True
    return False


# get the empirically pruned step for sampling


def get_epss_timesteps(n, device, dtype):
    dt = 1 / 32
    predefined_timesteps = {
        5: [0, 2, 4, 8, 16, 32],
        6: [0, 2, 4, 6, 8, 16, 32],
        7: [0, 2, 4, 6, 8, 16, 24, 32],
        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
    }
    t = predefined_timesteps.get(n, [])
    if not t:
        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
    return dt * torch.tensor(t, device=device, dtype=dtype)