From b62e0dc3c5574ce83a53720b9bd90acdff3eb399 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?= <41315874+fumiama@users.noreply.github.com> Date: Mon, 24 Jun 2024 16:27:15 +0900 Subject: [PATCH] optimize(utils): move custom processors into model (#419) --- ChatTTS/core.py | 2 +- ChatTTS/infer/api.py | 2 +- ChatTTS/model/gpt.py | 2 +- ChatTTS/model/processors.py | 45 +++++++++++++++++++++++++++++++++++++ ChatTTS/utils/gpu.py | 2 +- ChatTTS/utils/infer.py | 44 ------------------------------------ 6 files changed, 49 insertions(+), 48 deletions(-) create mode 100644 ChatTTS/model/processors.py diff --git a/ChatTTS/core.py b/ChatTTS/core.py index b1c5d1f..fb43a9a 100644 --- a/ChatTTS/core.py +++ b/ChatTTS/core.py @@ -119,7 +119,7 @@ class Chat: coef: Optional[str] = None ): if device is None: - device = select_device(4096) + device = select_device() self.logger.log(logging.INFO, f'use {device}') self.device = device diff --git a/ChatTTS/infer/api.py b/ChatTTS/infer/api.py index 5adb93b..94c3ad0 100644 --- a/ChatTTS/infer/api.py +++ b/ChatTTS/infer/api.py @@ -3,7 +3,7 @@ import torch import torch.nn.functional as F from transformers.generation import TopKLogitsWarper, TopPLogitsWarper -from ..utils.infer import CustomRepetitionPenaltyLogitsProcessorRepeat +from ..model.processors import CustomRepetitionPenaltyLogitsProcessorRepeat from ..utils.io import del_all from ..model.gpt import GPT diff --git a/ChatTTS/model/gpt.py b/ChatTTS/model/gpt.py index 32bb952..823d1f0 100644 --- a/ChatTTS/model/gpt.py +++ b/ChatTTS/model/gpt.py @@ -16,7 +16,7 @@ from transformers import LlamaModel, LlamaConfig, LogitsWarper from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPast -from ..utils.infer import CustomRepetitionPenaltyLogitsProcessorRepeat +from .processors import CustomRepetitionPenaltyLogitsProcessorRepeat from ..utils.io import del_all diff --git a/ChatTTS/model/processors.py b/ChatTTS/model/processors.py new file mode 100644 index 0000000..0c07cbe --- /dev/null +++ b/ChatTTS/model/processors.py @@ -0,0 +1,45 @@ +import torch +import torch.nn.functional as F + + +class CustomRepetitionPenaltyLogitsProcessorRepeat(): + + def __init__(self, penalty: float, max_input_ids, past_window): + if not isinstance(penalty, float) or not (penalty > 0): + raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") + + self.penalty = penalty + self.max_input_ids = max_input_ids + self.past_window = past_window + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + + input_ids = input_ids[:, -self.past_window:] + freq = F.one_hot(input_ids, scores.size(1)).sum(1) + freq[self.max_input_ids:] = 0 + alpha = self.penalty**freq + scores = scores.contiguous() + scores = torch.where(scores < 0, scores*alpha, scores/alpha) + + return scores + +class CustomRepetitionPenaltyLogitsProcessor(): + + def __init__(self, penalty: float, max_input_ids, past_window): + if not isinstance(penalty, float) or not (penalty > 0): + raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") + + self.penalty = penalty + self.max_input_ids = max_input_ids + self.past_window = past_window + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + + input_ids = input_ids[:, -self.past_window:] + score = torch.gather(scores, 1, input_ids) + _score = score.detach().clone() + score = torch.where(score < 0, score * self.penalty, score / self.penalty) + score[input_ids>=self.max_input_ids] = _score[input_ids>=self.max_input_ids] + scores.scatter_(1, input_ids, score) + + return scores diff --git a/ChatTTS/utils/gpu.py b/ChatTTS/utils/gpu.py index 9d3836e..248dff3 100644 --- a/ChatTTS/utils/gpu.py +++ b/ChatTTS/utils/gpu.py @@ -3,7 +3,7 @@ import torch from .log import logger -def select_device(min_memory=2048): +def select_device(min_memory=2047): if torch.cuda.is_available(): available_gpus = [] for i in range(torch.cuda.device_count()): diff --git a/ChatTTS/utils/infer.py b/ChatTTS/utils/infer.py index 63fcf9b..1288ebe 100644 --- a/ChatTTS/utils/infer.py +++ b/ChatTTS/utils/infer.py @@ -5,51 +5,7 @@ import sys from numba import jit import numpy as np -import torch -import torch.nn.functional as F - -class CustomRepetitionPenaltyLogitsProcessorRepeat(): - - def __init__(self, penalty: float, max_input_ids, past_window): - if not isinstance(penalty, float) or not (penalty > 0): - raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") - - self.penalty = penalty - self.max_input_ids = max_input_ids - self.past_window = past_window - - def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: - - input_ids = input_ids[:, -self.past_window:] - freq = F.one_hot(input_ids, scores.size(1)).sum(1) - freq[self.max_input_ids:] = 0 - alpha = self.penalty**freq - scores = scores.contiguous() - scores = torch.where(scores < 0, scores*alpha, scores/alpha) - - return scores - -class CustomRepetitionPenaltyLogitsProcessor(): - - def __init__(self, penalty: float, max_input_ids, past_window): - if not isinstance(penalty, float) or not (penalty > 0): - raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") - - self.penalty = penalty - self.max_input_ids = max_input_ids - self.past_window = past_window - - def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: - - input_ids = input_ids[:, -self.past_window:] - score = torch.gather(scores, 1, input_ids) - _score = score.detach().clone() - score = torch.where(score < 0, score * self.penalty, score / self.penalty) - score[input_ids>=self.max_input_ids] = _score[input_ids>=self.max_input_ids] - scores.scatter_(1, input_ids, score) - - return scores @jit def _find_index(table: np.ndarray, val: np.uint16):