libs/utils.py - master - diger/fb2f5tts

libs/utils.py
355 строк12 KB
DiGer
fix audio ref from gr.Audio
08 янв 2026, 14:32
08 янв 2026, 14:32ec93bbf
import os
import re
import json
import requests
import zipfile
import torch
import torchaudio
import gc
import numpy as np
import librosa
from vocos import Vocos
from ruaccent import RUAccent
from PIL import Image, ImageDraw, ImageFont
from libs.f5_tts_backend import load_models,generate_audio,_preloaded_speakers_data

device = 'CPU'

if torch.cuda.is_available():
    device = 'CUDA'

print(f'Run with {device}')

now_dir = os.getcwd()
data_path = os.path.join(now_dir, "data")
ab_name = ''

if not os.path.exists(data_path):
    os.makedirs(data_path)

with open('dict/word_dict.json', 'r', encoding="utf-8") as file:
    word_dict = json.load(file)

with open('dict/num_dict.json', 'r', encoding="utf-8") as fl:
    num_dict = json.load(fl)

args=''

class TTSModel:
    def __init__(self):
        self.vocoder = None
        self.model = None
        self.ver = None

    def load_vocoder(self):
        local_folder = "models"
        model_path = os.path.join(now_dir, local_folder, 'vocos-mel-24khz')
        model_filepath = os.path.join(now_dir, local_folder, 'vocos-mel-24khz.zip')
        if not os.path.isfile(model_filepath):
            print(f'Download vocoder')
            model_url = "https://myfreenet.ru/models/vocos-mel-24khz.zip"
            m, status = download_model(model_url,model_filepath)
            if m is None:
                    return m, status
            with zipfile.ZipFile(model_filepath, 'r') as zip_ref:
                    zip_ref.extractall(local_folder)
        vocoder = Vocos.from_hparams(f"{model_path}/config.yaml")
        state_dict = torch.load(f"{model_path}/pytorch_model.bin", map_location="cpu")
        vocoder.load_state_dict(state_dict)
        print('Vocoder loaded')
        self.vocoder = vocoder.eval().to(device.lower())

    def load(self, ver):
        local_folder = "models"
        self.ver = ver
        model_pt = 'model_last_inference.safetensors'
        if ver == 1:
            model_pt = 'espeech_tts_rlv2.pt'
        model_filepath = os.path.join(now_dir, local_folder, model_pt)
        if not os.path.isfile(model_filepath):
            print(f'Download {model_pt}')
            model_url = f"https://myfreenet.ru/models/{model_pt}"
            m, status = download_model(model_url,model_filepath)
            if m is None:
                return m, status
   
        if self.model is not None:
            del self.model
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            gc.collect()
        self.model = load_models(ckpt_path=model_filepath)
        return ver, "Модель успешно загружена!"
    
    def synth_audio(self, text, speaker_id=0, speed=0.9, ref_audio=None, ref_text=''):
        audio_wave, sample_rate = generate_audio(
            text,
            model_obj=self.model,
            vocoder=self.vocoder,
            speed=speed,
            ref_data = get_speaker_data(speaker_id, custom_audio=ref_audio, text=ref_text)
        )
        audio_wave = (audio_wave * 32767).astype(np.int16)
        return audio_wave, sample_rate

synth = TTSModel()

class ACCModel:
    def __init__(self):
        self.accentizer = None
        self.ver = None

    def load(self, ver):
        self.ver = ver
        if ver == 1:
            self.accentizer = RUAccent()
            self.accentizer.load(
                omograph_model_size='big_poetry',
                use_dictionary=True,
                device=device,
                workdir="./models"
            )
            return ver, "Модель успешно загружена!"
        else:
            silero_stress = 'accentor.pt'
            silero_directory = 'models/silero_stress'
            silero_filepath = os.path.join(now_dir, silero_directory, silero_stress)
            if not os.path.isfile(silero_filepath):
                os.makedirs(silero_directory, exist_ok=True)
                print(f'Download silero stress')
                model_url = "https://github.com/snakers4/silero-stress/raw/refs/heads/master/src/silero_stress/data/accentor.pt"
                m, status = download_model(model_url,silero_filepath)
                if m is None:
                    return m, status
            self.accentizer = torch.package.PackageImporter(silero_filepath).load_pickle("accentor_models", "accentor")
            quantized_weight = self.accentizer.homosolver.model.bert.embeddings.word_embeddings.weight.data.clone()
            restored_weights = self.accentizer.homosolver.model.bert.scale * (quantized_weight - self.accentizer.homosolver.model.bert.zero_point)
            self.accentizer.homosolver.model.bert.embeddings.word_embeddings.weight.data = restored_weights
            return ver, "Модель успешно загружена!"

    def process_accent(self, string, regexp):
        if self.ver == 1:
            return self.accentizer.process_all(string, regexp)
        if not regexp:
            return self.accentizer(string)
        
        pattern = re.compile(regexp)
        matches = list(pattern.finditer(string))
        
        if not matches:
            return self.accentizer(string)

        result_parts = []
        prev_end = 0
        
        for match in matches:
            start, end = match.start(), match.end()
            result_parts.append(self.accentizer(string[prev_end:start]))
            result_parts.append(string[start:end])
            
            prev_end = end
        
        result_parts.append(self.accentizer(string[prev_end:]))
        
        return "".join(result_parts)

accentizer = ACCModel()


def download_model(model_url, target_path):
    try:
        response = requests.get(model_url, stream=True, timeout=5)
        response.raise_for_status()
        expected_size = int(response.headers.get('content-length', 0))
        with open(target_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        actual_size = os.path.getsize(target_path)
        if actual_size > 0 and (expected_size == 0 or actual_size == expected_size):
            return True, True
        else:
            if os.path.exists(target_path):
                os.remove(target_path)
            return None, 'Error: Размер неверный!'

    except Exception as e:
        if os.path.exists(target_path):
            os.remove(target_path)
        return None, f'Error: {e}'

def get_spk_list():
    spk_list = []
    for i in _preloaded_speakers_data:
        spk_list.append((_preloaded_speakers_data[i]['name'],i))
    return spk_list

def set_args(s_args):
    global args
    args = s_args
    return args

def get_args():
    return args

def convert(seconds):
    min, sec = divmod(seconds, 60)
    min = f'{int(min)} m. ' if min else ''
    sec = f'{int(sec)} s.'

    return min + sec

def set_ab_name(d_path):
    global ab_name
    ab_name = d_path
    return ab_name

def get_ab_name():
    return ab_name

def get_data_list(d_path=data_path):
    return [
        dirpath
        for dirpath in os.listdir(d_path)
    ]

def convert_to_jpg(image,dest_image):
    img = Image.fromarray(image)
    img.save(dest_image)
    return img

def load_image(ab_name):
    d_path = os.path.join(now_dir, "data", ab_name)
    destination_path = os.path.join(d_path, 'cover.jpg')
    if not os.path.exists(destination_path):
        destination_path = os.path.join(now_dir, 'libs', "cover.jpg")
    img = Image.open(destination_path)
    return img

def add_text_cover(output_path, autor, title):
    image = Image.open('libs/cover.jpg')
    draw = ImageDraw.Draw(image)
    font1 = ImageFont.truetype('libs/Horovod-Regular.ttf', size=22)
    font2 = ImageFont.truetype('libs/Horovod-Regular.ttf', size=26)
    
    bbox = draw.textbbox((0, 0), autor, font=font1)
    text_width, text_height = bbox[2] - bbox[0], bbox[3] - bbox[1]
    x = (image.width - text_width) / 2
    y = (image.height - text_height) / 3
    draw.text((x, y), autor, font=font1, fill='black')
    
    bbox = draw.textbbox((0, 0), title, font=font2)
    text_width, text_height = bbox[2] - bbox[0], bbox[3] - bbox[1]
    x = (image.width - text_width) / 2
    y = (image.height - text_height) / 3 + 24
    draw.text((x, y), title, font=font2, fill='black')

    image.save(output_path)

def change_pitch(audio_array, sample_rate, semitones):
    audio_array = audio_array.astype(np.float32) / 32768.0

    if np.max(np.abs(audio_array)) > 1.0:
        audio_array = audio_array / np.max(np.abs(audio_array))
    actual_semitones = (semitones - 50) / 50 * 12
    y_shifted = librosa.effects.pitch_shift(
        y=audio_array,
        sr=sample_rate,
        n_steps=actual_semitones
    )

    y_shifted_int16 = (y_shifted * 32767.0).astype(np.int16)
    return y_shifted_int16

def prep_audio(
    ref_audio,
    target_sample_rate: int = 24000,
    device: str = "cuda" if torch.cuda.is_available() else "cpu",
):
    audio = None
    if type(ref_audio) is tuple:
        sr, waveform_data = ref_audio
        audio = torch.from_numpy(waveform_data)
        audio = torch.as_tensor(audio, dtype=torch.float32)
        a_len = audio.shape[-1] // 256
        
        if sr != target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
            audio = resampler(audio)
            sr = target_sample_rate
        target_rms = 0.1
        rms = torch.sqrt(torch.mean(torch.square(audio)))
        if rms < target_rms:
            audio = audio * target_rms / rms
            rms = torch.sqrt(torch.mean(torch.square(audio)))

        rms_numpy = rms.cpu().numpy()
        rms_tensor = torch.from_numpy(rms_numpy).to(device)
        
        return audio.to(device), rms_tensor, target_sample_rate, a_len
    else:
        audio, sr = torchaudio.load(ref_audio)
        if audio.shape[0] > 1:
            audio = torch.mean(audio, dim=0, keepdim=True)
        
        a_len = audio.shape[-1] // 256
        # Нормализация громкости
        target_rms = 0.1
        rms = torch.sqrt(torch.mean(torch.square(audio)))
        if rms < target_rms:
            audio = audio * target_rms / rms
            rms = torch.sqrt(torch.mean(torch.square(audio)))
        
        return audio.to(device), rms, target_sample_rate, a_len

def prep_audio_from_gradio(audio_input, target_sample_rate=24000, device=None):

    if isinstance(audio_input, tuple):
        sr, audio_array = audio_input
        if audio_array.dtype == np.int16:
            audio_array = audio_array / 32768.0
        audio = torch.from_numpy(audio_array).float().unsqueeze(0)  # (1, T)
    else:
        raise ValueError(f"Ожидался tuple (sr, np.ndarray), получено: {type(audio_input)}")

    if audio.shape[0] > 1:
        audio = torch.mean(audio, dim=0, keepdim=True)

    if sr != target_sample_rate:
        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
        audio = resampler(audio)
    else:
        audio = audio

    target_rms_val = 0.1
    rms = torch.sqrt(torch.mean(torch.square(audio)))
    if rms < target_rms_val:
        audio = audio * target_rms_val / rms
        rms = torch.tensor(target_rms_val)

    audio_len = audio.shape[-1] // 256

    return {
        'audio': audio.to(device),
        'rms': rms,
        'audio_len': audio_len
    }

def get_speaker_data(speaker_id=None, custom_audio=None, text=""):
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if custom_audio is not None:
        # Используем кастомное аудио
        data = prep_audio_from_gradio(custom_audio, device=device)
        data['text'] = text
        data['text_len'] = len(text.encode("utf-8"))
    elif speaker_id is not None:
        # Используем спикера из архива
        speaker_key = str(speaker_id)
        if speaker_key not in _preloaded_speakers_data:
            raise ValueError(f"Speaker {speaker_id} not found in archive")
        data = _preloaded_speakers_data[speaker_key].copy()
    else:
        raise ValueError("Должен быть указан speaker_id или custom_audio")
    
    return data