From 2f5669650651fc79b81e2d6bbaa895c847f570f3 Mon Sep 17 00:00:00 2001 From: mrq Date: Mon, 11 Nov 2024 20:21:16 -0600 Subject: [PATCH] overhauled inference/sampler kwargs to stop being a bloated mess --- vall_e/__main__.py | 53 ++-- vall_e/config.py | 44 +-- vall_e/data.py | 19 +- vall_e/inference.py | 122 +-------- vall_e/models/ar_nar.py | 568 +++++++++++++++------------------------ vall_e/models/base.py | 46 ++-- vall_e/utils/__init__.py | 3 +- vall_e/utils/utils.py | 16 ++ vall_e/webui.py | 198 +++++++------- 9 files changed, 431 insertions(+), 638 deletions(-) diff --git a/vall_e/__main__.py b/vall_e/__main__.py index 96d7508..d51e0fd 100755 --- a/vall_e/__main__.py +++ b/vall_e/__main__.py @@ -20,20 +20,23 @@ def main(): parser.add_argument("--model", type=Path, default=None) parser.add_argument("--lora", type=Path, default=None) - parser.add_argument("--max-ar-steps", type=int, default=12 * cfg.dataset.frames_per_second) - parser.add_argument("--max-nar-levels", type=int, default=7) + parser.add_argument("--max-duration", type=int, default=12 * cfg.dataset.frames_per_second) + parser.add_argument("--max-steps", type=int, default=25) + parser.add_argument("--max-levels", type=int, default=7) - parser.add_argument("--ar-temp", type=float, default=0.5) - parser.add_argument("--nar-temp", type=float, default=0.0) - parser.add_argument("--min-ar-temp", type=float, default=-1.0) - parser.add_argument("--min-nar-temp", type=float, default=-1.0) + parser.add_argument("--ar-temperature", type=float, default=1.0) + parser.add_argument("--nar-temperature", type=float, default=0.0) + parser.add_argument("--min-ar-temperature", type=float, default=-1.0) + parser.add_argument("--min-nar-temperature", type=float, default=-1.0) parser.add_argument("--input-prompt-length", type=float, default=3.0) parser.add_argument("--input-prompt-prefix", action="store_true") + parser.add_argument("--prefix-silence", type=float, default=0.0) + parser.add_argument("--cfg-strength", type=float, default=0.0) parser.add_argument("--top-p", type=float, default=1.0) parser.add_argument("--top-k", type=int, default=0) parser.add_argument("--min-p", type=float, default=0.0) - parser.add_argument("--repetition-penalty", type=float, default=1.5) + parser.add_argument("--repetition-penalty", type=float, default=1.0) parser.add_argument("--repetition-penalty-decay", type=float, default=0.0) parser.add_argument("--length-penalty", type=float, default=0.0) parser.add_argument("--beam-width", type=int, default=0) @@ -73,17 +76,13 @@ def main(): config = args.model tts = TTS( config=config, lora=args.lora, device=args.device, dtype=args.dtype, amp=args.amp, attention=args.attention ) - output = tts.inference( - text=args.text, - references=args.references, - language=args.language, - task=args.task, - out_path=args.out_path, - input_prompt_length=args.input_prompt_length, - input_prompt_prefix=args.input_prompt_prefix, - max_ar_steps=args.max_ar_steps, max_nar_levels=args.max_nar_levels, - ar_temp=args.ar_temp, nar_temp=args.nar_temp, - min_ar_temp=args.min_ar_temp, min_nar_temp=args.min_nar_temp, + + sampling_kwargs = dict( + max_steps=args.max_steps, + max_levels=args.max_levels, + max_duration=args.max_duration, + ar_temperature=args.ar_temperature, nar_temperature=args.nar_temperature, + min_ar_temperature=args.min_ar_temperature, min_nar_temperature=args.min_nar_temperature, top_p=args.top_p, top_k=args.top_k, min_p=args.min_p, repetition_penalty=args.repetition_penalty, repetition_penalty_decay=args.repetition_penalty_decay, length_penalty=args.length_penalty, @@ -96,9 +95,23 @@ def main(): layer_skip_entropy_threshold=args.layer_skip_entropy_threshold, layer_skip_varentropy_threshold=args.layer_skip_varentropy_threshold, refine_on_stop=args.refine_on_stop, - - load_from_artifact=args.load_from_artifact, denoise_start=args.denoise_start, + input_prompt_prefix=args.input_prompt_prefix, + prefix_silence=args.prefix_silence, + cfg_strength=args.cfg_strength, + ) + + output = tts.inference( + text=args.text, + references=args.references, + language=args.language, + task=args.task, + out_path=args.out_path, + + input_prompt_length=args.input_prompt_length, + load_from_artifact=args.load_from_artifact, + + sampling_kwargs=sampling_kwargs, seed=args.seed, ) diff --git a/vall_e/config.py b/vall_e/config.py index ca570ba..770dbf1 100755 --- a/vall_e/config.py +++ b/vall_e/config.py @@ -292,6 +292,7 @@ class Model: #loss_factors: dict = field(default_factory=lambda: { "text": 0.1, "prom": 1.0, "resp": 1.0 }) # disable it by default since it causes a little more harm than good loss_factors: dict = field(default_factory=lambda: {}) capabilities: list = field(default_factory=lambda: ["ar", "nar"]) # + ["lang", "tone"] if you have your dataset labeled for such + kwargs: dict = field(default_factory=lambda: {}) experimental: dict | ModelExperimentalSettings | None = None # experimental settings @@ -410,6 +411,11 @@ class Model: return dict(include=include, exclude=exclude) + # to-do: derive default arguments from here + @property + def get_kwargs(self, type): + return self.kwargs + # should be renamed to Adapters @dataclass() class LoRA: @@ -466,32 +472,30 @@ class Evaluation: # necessary in order to make it not confusing with requiring not-directyl exposed arguments passed to the model @cached_property def ar_kwargs( self ): - kwargs = {} | self.kwargs return dict( - max_steps=kwargs.pop("max_ar_steps", 500), - sampling_temperature=kwargs.pop("ar_temp", 0.5), - sampling_min_temperature=kwargs.pop("min_ar_temp", -1), - sampling_top_p=kwargs.pop("top_p", 1.0), sampling_top_k=kwargs.pop("top_k", 0), sampling_min_p=kwargs.pop("min_p", 0.0), - sampling_repetition_penalty=kwargs.pop("repetition_penalty", 1.125), sampling_repetition_penalty_decay=kwargs.pop("repetition_penalty_decay", 0), - sampling_length_penalty=kwargs.pop("length_penalty", 0), - sampling_beam_width=kwargs.pop("beam_width", 0), - sampling_mirostat_tau=kwargs.pop("mirostat_tau", 0), - sampling_mirostat_eta=kwargs.pop("mirostat_eta", 0), - sampling_dry_multiplier=kwargs.pop("dry_multiplier", 0), - sampling_dry_base=kwargs.pop("dry_base", 0), - sampling_dry_allowed_length=kwargs.pop("dry_allowed_length", 0), - sampling_entropix=kwargs.pop("entropix_sampling", False), + max_steps=self.kwargs.get("max_ar_steps", 500), + temperature=self.kwargs.get("ar_temperature", 1.0), + min_temperature=self.kwargs.get("min_ar_temperature", -1), + top_p=self.kwargs.get("top_p", 1.0), top_k=self.kwargs.get("top_k", 0), min_p=self.kwargs.get("min_p", 0.0), + repetition_penalty=self.kwargs.get("repetition_penalty", 1.0), repetition_penalty_decay=self.kwargs.get("repetition_penalty_decay", 0), + length_penalty=self.kwargs.get("length_penalty", 0), + beam_width=self.kwargs.get("beam_width", 0), + mirostat_tau=self.kwargs.get("mirostat_tau", 0), + mirostat_eta=self.kwargs.get("mirostat_eta", 0), + dry_multiplier=self.kwargs.get("dry_multiplier", 0), + dry_base=self.kwargs.get("dry_base", 0), + dry_allowed_length=self.kwargs.get("dry_allowed_length", 0), + entropix=self.kwargs.get("entropix_sampling", False), ) @cached_property def nar_kwargs( self ): - kwargs = {} | self.kwargs return dict( - max_levels=kwargs.pop("max_nar_levels", 0), - sampling_temperature=kwargs.pop("nar_temp", 0.0), - sampling_min_temperature=kwargs.pop("min_nar_temp", -1), - sampling_top_p=kwargs.pop("top_p", 1.0), sampling_top_k=kwargs.pop("top_k", 0.0), sampling_min_p=kwargs.pop("min_p", 0.0), - sampling_repetition_penalty=kwargs.pop("repetition_penalty", 1.0), sampling_repetition_penalty_decay=kwargs.pop("repetition_penalty_decay", 0.0), + max_levels=self.kwargs.get("max_nar_levels", 0), + temperature=self.kwargs.get("nar_temperature", 0.0), + min_temperature=self.kwargs.get("min_nar_temp", -1), + top_p=self.kwargs.get("top_p", 1.0), top_k=self.kwargs.get("top_k", 0.0), min_p=self.kwargs.get("min_p", 0.0), + repetition_penalty=self.kwargs.get("repetition_penalty", 1.0), repetition_penalty_decay=self.kwargs.get("repetition_penalty_decay", 0.0), ) @dataclass() diff --git a/vall_e/data.py b/vall_e/data.py index 72e85ff..7485e70 100755 --- a/vall_e/data.py +++ b/vall_e/data.py @@ -571,7 +571,7 @@ def _load_paths_from_metadata(group_name, type="training", validate=False): _fn = _get_hdf5_paths if cfg.dataset.use_hdf5 else _get_paths_of_extensions def key( id, entry=None ): - return f"/{type}/{_get_hdf5_path(data_dir)}/{id}" if cfg.dataset.use_hdf5 else data_dir / id + return f"/{type}/{_get_hdf5_path(data_dir)}/{id}" if cfg.dataset.use_hdf5 else str(data_dir / id) metadata_path = cfg.metadata_dir / f'{group_name}.json' metadata = {} @@ -628,21 +628,8 @@ def _get_hdf5_paths( data_dir, type="training", validate=False ): def _get_paths_of_extensions( path, extensions=_get_quant_extension(), validate=False ): if isinstance(path, str): path = Path(path) - - def _validate(path): - if "".join(path.suffixes) not in extensions: - return False - if not _get_phone_path(path).exists() or not _get_quant_path(path).exists(): - return False - if not validate: - return True - # to-do: find an easy way to determine size from pickled quants without loading - # to-do: find a consistent way to derive phoneme count from filesize (probably can't due to utf-8) - phones = len(_get_phones(_get_phone_path(path))) # _get_phone_path(path).stat().st_size // 2 + 1 - return cfg.dataset.min_phones <= phones and phones <= cfg.dataset.max_phones - - - return [ p for p in list(path.iterdir()) if _validate(p) ] if path.exists() and path.is_dir() else [] + + return [ p for p in list(path.iterdir()) ] if path.exists() and path.is_dir() else [] def _load_quants(path, return_metadata=False) -> Tensor: qnt = np.load(_get_quant_path(path), allow_pickle=True)[()] diff --git a/vall_e/inference.py b/vall_e/inference.py index 8a39d8a..5929118 100755 --- a/vall_e/inference.py +++ b/vall_e/inference.py @@ -186,55 +186,15 @@ class TTS(): references, language="en", task="tts", - # - max_ar_steps=6 * cfg.dataset.frames_per_second, - max_nar_levels=7, - # - input_prompt_length=0.0, - input_prompt_prefix=False, - prefix_silence=0.0, - # - ar_temp=0.0, - nar_temp=0.0, - # - min_ar_temp=0.0, - min_nar_temp=0.0, - # - top_p=1.0, - top_k=0, - min_p=0.0, - # - repetition_penalty=1.0, - repetition_penalty_decay=0.0, - length_penalty=0.0, - # - beam_width=0, - # - mirostat_tau=0, - mirostat_eta=0.1, - # - dry_multiplier=0.0, - dry_base=1.75, - dry_allowed_length=2, - # - entropix_sampling=False, - # - layer_skip=False, - layer_skip_exit_layer=-1, - layer_skip_entropy_threshold=-1, - layer_skip_varentropy_threshold=-1, - # - refine_on_stop=False, - # + + input_prompt_length = 0, + load_from_artifact = False, + seed = None, - # - load_from_artifact = None, - denoise_start = 0.0, - out_path=None, - tqdm=True, use_lora=None, + **sampling_kwargs, ): lines = text.split("\n") @@ -265,25 +225,10 @@ class TTS(): with torch.autocast("cuda", dtype=self.dtype, enabled=self.amp): if model_ar is not None: text_list = model_ar( - text_list=None, proms_list=[resp], lang_list=[lang], resps_list=[resp], max_steps=max_ar_steps, task_list=["stt"], - sampling_temperature=ar_temp, - sampling_min_temperature=min_ar_temp, - sampling_top_p=top_p, sampling_top_k=top_k, sampling_min_p=min_p, - sampling_repetition_penalty=repetition_penalty, sampling_repetition_penalty_decay=repetition_penalty_decay, - sampling_length_penalty=length_penalty, - sampling_beam_width=beam_width, - sampling_mirostat_tau=mirostat_tau, - sampling_mirostat_eta=mirostat_eta, - sampling_dry_multiplier=dry_multiplier, - sampling_dry_base=dry_base, - sampling_dry_allowed_length=dry_allowed_length, - sampling_entropix=entropix_sampling, - sampling_layer_skip=layer_skip, - sampling_layer_skip_exit_layer=layer_skip_exit_layer, - sampling_refine_on_stop=refine_on_stop, - + text_list=None, proms_list=[resp], lang_list=[lang], resps_list=[resp], task_list=["stt"], disable_tqdm=not tqdm, use_lora=use_lora, + **sampling_kwargs, ) else: raise Exception("!") @@ -292,10 +237,6 @@ class TTS(): return text_list[0] - # validate settings here - if not references and ar_temp < 0.5: - _logger.warning(f'Audio-promptless inferencing fails with low AR temperatures.') - for line in lines: if out_path is None: output_dir = Path("./data/results/") @@ -315,52 +256,21 @@ class TTS(): with torch.autocast("cuda", dtype=self.dtype, enabled=self.amp): if model_ar is not None: resps_list = model_ar( - text_list=[phns], proms_list=[prom], lang_list=[lang], max_steps=max_ar_steps, task_list=["tts"], - input_prompt_prefix=input_prompt_prefix, - prefix_silence=prefix_silence, - sampling_temperature=ar_temp, - sampling_min_temperature=min_ar_temp, - sampling_top_p=top_p, sampling_top_k=top_k, sampling_min_p=min_p, - sampling_repetition_penalty=repetition_penalty, sampling_repetition_penalty_decay=repetition_penalty_decay, - sampling_length_penalty=length_penalty, - sampling_beam_width=beam_width, - sampling_mirostat_tau=mirostat_tau, - sampling_mirostat_eta=mirostat_eta, - sampling_dry_multiplier=dry_multiplier, - sampling_dry_base=dry_base, - sampling_dry_allowed_length=dry_allowed_length, - sampling_entropix=entropix_sampling, - sampling_layer_skip=layer_skip, - sampling_layer_skip_exit_layer=layer_skip_exit_layer, - sampling_layer_skip_entropy_threshold=layer_skip_entropy_threshold, - sampling_layer_skip_varentropy_threshold=layer_skip_varentropy_threshold, - sampling_refine_on_stop=refine_on_stop, - + text_list=[phns], proms_list=[prom], lang_list=[lang], task_list=["tts"], disable_tqdm=not tqdm, use_lora=use_lora, + **sampling_kwargs, ) resps_list = model_nar( text_list=[phns], proms_list=[prom], lang_list=[lang], resps_list=resps_list, task_list=["tts"], - input_prompt_prefix=input_prompt_prefix, - max_levels=max_nar_levels, - sampling_temperature=nar_temp, - sampling_min_temperature=min_nar_temp, - sampling_top_p=top_p, sampling_top_k=top_k, sampling_min_p=min_p, - sampling_repetition_penalty=repetition_penalty, sampling_repetition_penalty_decay=repetition_penalty_decay, - sampling_layer_skip=layer_skip, - sampling_layer_skip_exit_layer=layer_skip_exit_layer, - sampling_layer_skip_entropy_threshold=layer_skip_entropy_threshold, - sampling_layer_skip_varentropy_threshold=layer_skip_varentropy_threshold, - disable_tqdm=not tqdm, use_lora=use_lora, + **sampling_kwargs, ) elif model_len is not None: - len_list = model_len( text_list=[phns], proms_list=[prom], task_list=["len"], max_steps=5, disable_tqdm=not tqdm ) # don't need more than that - len_list = [ clamp(l, 1, max_ar_steps) for l in len_list ] + len_list = model_len( text_list=[phns], proms_list=[prom], task_list=["len"], disable_tqdm=not tqdm, **{"max_steps": 5} ) # don't need more than that kwargs = {} - # nasty hardcode to load a reference file and have that as the input target if load_from_artifact and load_from_artifact.exists(): artifact = np.load(load_from_artifact, allow_pickle=True)[()] @@ -373,17 +283,9 @@ class TTS(): kwargs["resps_list"] = [ resp[:, :1] ] resps_list = model_nar( text_list=[phns], proms_list=[prom], len_list=len_list, task_list=["tts"], - max_steps=max_ar_steps, - max_levels=max_nar_levels, - sampling_temperature=nar_temp, - sampling_min_temperature=min_nar_temp, - sampling_top_p=top_p, sampling_top_k=top_k, sampling_min_p=min_p, - sampling_repetition_penalty=repetition_penalty, sampling_repetition_penalty_decay=repetition_penalty_decay, - denoise_start=denoise_start, - disable_tqdm=not tqdm, use_lora=use_lora, - **kwargs, + **(sampling_kwargs | kwargs), ) else: raise Exception("!") diff --git a/vall_e/models/ar_nar.py b/vall_e/models/ar_nar.py index e898b61..d9cc325 100644 --- a/vall_e/models/ar_nar.py +++ b/vall_e/models/ar_nar.py @@ -17,14 +17,14 @@ import math import time from einops import rearrange from torch import Tensor -from tqdm import trange +from tqdm import trange, tqdm import logging _logger = logging.getLogger(__name__) from ..emb.qnt import trim, encode_as_embedding, get_silence -from ..utils import get_devices, setup_logging, timer, clamp +from ..utils import get_devices, setup_logging, timer, clamp, convert_kwargs from .lora import enable_lora @@ -187,6 +187,149 @@ class AR_NAR(Base): quant_levels=quant_levels, ) + def forward_nar_masked( + self, + + text_list: list[Tensor], + proms_list: list[Tensor], + resps_list: list[Tensor] | None = None, + + task_list: list[Tensor] | None = None, + lang_list: list[Tensor] | None = None, + tone_list: list[Tensor] | None = None, + len_list: list[Tensor] | None = None, + + disable_tqdm=False, + use_lora=None, + **sampling_kwargs, + ): + device = text_list[0].device + batch_size = len(text_list) + + # special "scheduling" to inference RVQ-level 0 + level = 0 + if cfg.lora is not None: + enable_lora( self, cfg.lora.active_level( level ) if use_lora is None else use_lora ) + + def log(x, eps = 1e-20): + return torch.log(x.clamp(min = eps)) + + def gumbel_sample(x, temperature = 1., dim = -1): + return ((x / max(temperature, 1e-10)) + -log(-log(torch.zeros_like(x).uniform_(0, 1)))).argmax(dim = dim) + + # convert (N)AR specific args + sampling_kwargs = convert_kwargs( sampling_kwargs, "ar_" ) + + max_length = sampling_kwargs.pop("max_duration", 500) + max_steps = sampling_kwargs.get("max_steps", 25) + + temperature = sampling_kwargs.pop("temperature", 1.0) + cfg_strength = sampling_kwargs.get("cfg_strength", 0.0) + start_noise = sampling_kwargs.get("denoise_start", 0.0) + end_noise = sampling_kwargs.get("denoise_end", 1.0) + max_steps = math.floor(max_steps * (end_noise - start_noise)) + + len_list = [ clamp(l, 1, max_length) for l in len_list ] + + # if we're denoising from an existing sequence + if start_noise > 0.0 and resps_list is not None: + noise_p = math.cos( start_noise * math.pi * 0.5 ) + mask = [ torch.tensor( [ random.random() < noise_p for _ in range( seq_len ) ], dtype=torch.bool, device=device ) for seq_len in len_list ] + resps_list = [ torch.where( mask, self.stop_token, resps[:, 0] ) for seq_len, resps in zip( len_list, resps_list ) ] + else: + resps_list = [ torch.ones((seq_len,), dtype=torch.int16, device=device) * self.stop_token for seq_len in len_list ] + + scores = [ torch.zeros((seq_len,), dtype=torch.float32, device=device) for seq_len in len_list ] + quant_levels = [ level for _ in range(batch_size) ] + null_text = [ torch.tensor([1, 2], device=device, dtype=torch.int16) for _ in range(batch_size) ] + null_prom = [ None for _ in range(batch_size) ] + prev_list = resps_list + + for timestep, steps_until_x0 in tqdm(zip(torch.linspace(start_noise, end_noise, max_steps), reversed(range(max_steps))), desc="NAR Masked", disable=disable_tqdm, total=max_steps): + # get noise level, per cosine scheduling + noise_p = math.cos( timestep * math.pi * 0.5 ) + # pick the worst scoring tokens to mask off + masked_indices = [ score.topk( max(int( noise_p * seq_len ), 1), dim=-1 ).indices for score, seq_len in zip(scores, len_list) ] + # mask off inputs + resps_list = [ resp.scatter(0, indices, self.stop_token) for resp, indices in zip( resps_list, masked_indices ) ] + # boolean mask + is_masked = [ resps == self.stop_token for resps in resps_list ] + + time_list = [ timestep for _ in range(batch_size) ] + + # setup inputs + inputs = super().inputs( + text_list=text_list, + proms_list=proms_list, + resps_list=resps_list, + lang_list=lang_list, + tone_list=tone_list, + time_list=time_list, + quant_levels=quant_levels, + ) + output = super().forward( + inputs=inputs, + quant_levels=quant_levels, + #layer_skip_variables=sampling_layer_skip_variables, + ) + + logits = output.logits + + if cfg_strength > 0: + null_inputs = super().inputs( + text_list=null_text, + proms_list=null_prom, + resps_list=resps_list, + lang_list=lang_list, + tone_list=tone_list, + time_list=time_list, + quant_levels=quant_levels, + ) + null_output = super().forward( + inputs=null_inputs, + quant_levels=quant_levels, + #layer_skip_variables=sampling_layer_skip_variables, + ) + for seq_len, logit, null_logit in zip(len_list, output.logits, null_output.logits): + logit[-seq_len:] = null_logit[-seq_len:] + ( logit[-seq_len:] - null_logit[-seq_len:] ) * cfg_strength + + # sample with sampler settings + filtered_sampled = super().sample( + logits=logits, + prev_list=prev_list, + quant_levels=quant_levels, + + temperature=temperature * (steps_until_x0 / max_steps) , + **sampling_kwargs, + ) + + # retrieves unfiltered logits + unfiltered_sampled = super().sample( + logits=logits, + prev_list=prev_list, + quant_levels=quant_levels, + + temperature=0.0, + **sampling_kwargs, + ) + # update previous list of tokens + prev_list = resps_list + + # sample with gumbelnoise + # I actually feel like this doesn't matter? it's hard to judge with a partially trained NAR-len model + sampled_ids = [ gumbel_sample( logits, temperature=temperature, dim=-1 ) for logits in filtered_sampled.logits[0] ] + #sampled_ids = filtered_sampled[0] + + # keep unmasked tokens + resps_list = [ torch.where( masked, input_ids, resps ) for masked, input_ids, resps in zip( is_masked, sampled_ids, resps_list ) ] + # update scores (conjugated to put the worst scores at the top) + scores = [ 1.0 - torch.tensor([score for score in scores], device=device) for scores in unfiltered_sampled.scores ] + + if cfg.experimental and max_steps > 0: + print( timestep, steps_until_x0, noise_p, resps_list, scores ) + + return resps_list + def forward_nar( self, text_list: list[Tensor], @@ -198,40 +341,9 @@ class AR_NAR(Base): tone_list: list[Tensor] | None = None, len_list: list[Tensor] | None = None, - training: bool | int | None = None, - - max_steps: int = 1000, - max_levels: int = 0, - - input_prompt_prefix: bool = False, - prefix_silence: float = 1.0, - denoise_start: float = 0.0, - - sampling_temperature: float = 1.0, - sampling_min_temperature: float = -1.0, - sampling_top_k: int = -100, - sampling_top_p: float = 1.0, - sampling_min_p: float = 0.0, - sampling_repetition_penalty: float = 1.0, - sampling_repetition_penalty_decay: float = 0.0, - sampling_length_penalty: float = 0.0, - sampling_beam_width: int = 0, - sampling_mirostat_tau: float = 0.0, - sampling_mirostat_eta: float = 0.1, - sampling_dry_multiplier=0.0, - sampling_dry_base=1.75, - sampling_dry_allowed_length=2, - sampling_entropix=False, - - sampling_layer_skip: bool = False, - sampling_layer_skip_exit_layer: int = -1, - sampling_layer_skip_entropy_threshold: float = -1, - sampling_layer_skip_varentropy_threshold: float = -1, - - sampling_refine_on_stop: bool = False, - disable_tqdm=False, use_lora=None, + **sampling_kwargs, ): # deduce batch_size if text_list is not None: @@ -243,9 +355,15 @@ class AR_NAR(Base): device = resps_list[0].device batch_size = len(resps_list) + + max_levels = sampling_kwargs.get("max_levels", 0) + # convert NAR specific args + sampling_kwargs = convert_kwargs( sampling_kwargs, "nar_" ) + if max_levels == 0: max_levels = self.n_max_levels - 1 + """ sampling_layer_skip_variables = {} if sampling_layer_skip else None if sampling_layer_skip: @@ -255,162 +373,20 @@ class AR_NAR(Base): sampling_layer_skip_variables["varentropy_threshold"] = sampling_layer_skip_varentropy_threshold if sampling_layer_skip_exit_layer >= 0: sampling_layer_skip_variables["max_layer"] = sampling_layer_skip_exit_layer + """ # inference NAR level 0 if len_list is not None: - mask_token = torch.tensor([self.stop_token], dtype=torch.int16, device=device) - prev_list = [ torch.concat([ mask_token for _ in range( resp_len ) ]) for resp_len in len_list ] - - # special "scheduling" to inference RVQ-level 0 - level = 0 - if cfg.lora is not None: - enable_lora( self, cfg.lora.active_level( level ) if use_lora is None else use_lora ) - - def log(x, eps = 1e-20): - return torch.log(x.clamp(min = eps)) - - def gumbel_sample(x, temperature = 1., dim = -1): - return ((x / max(temperature, 1e-10)) + -log(-log(torch.zeros_like(x).uniform_(0, 1)))).argmax(dim = dim) - - _super = super() - # to-do: allow for batch processing (it should probably work batched anyways) - def demask_sampling( batch_index, seq_len ): - # overrides, to be user-controllable soonTM - max_steps = 10 - temperature = 1.0 - cfg_strength = 1.0 - sampling_repetition_penalty = 1.0 # force rep pen off, because this caused false positives due to how rep pen was being naively applied...... - sampling_top_p = 0.9 # a lot of demasking samplers use a top-k of seq_len * 0.9 - - start_temperature = temperature - start_noise = 0.0 - end_noise = 1.0 - - # if we're denoising from an existing sequence - if denoise_start > 0.0 and resps_list is not None: - start_noise = denoise_start - noise_p = math.cos( start_noise * math.pi * 0.5 ) - mask = torch.tensor( [ random.random() < noise_p for _ in range( seq_len ) ], dtype=torch.bool, device=device ) - input_ids = torch.where( mask, self.stop_token, resps_list[batch_index][:, 0] ) - else: - input_ids = torch.ones((seq_len,), dtype=torch.int16, device=device) * self.stop_token - - scores = torch.zeros((seq_len,), dtype=torch.float32, device=device) - - quant_levels = [ level for _ in range(batch_size) ] - prev_list = [ input_ids ] - - - null_text = torch.tensor([1, 2], device=device, dtype=torch.int16) - null_prom = None - - max_steps = math.floor(max_steps * (end_noise - start_noise)) - - for timestep, steps_until_x0 in zip(torch.linspace(start_noise, end_noise, max_steps), reversed(range(max_steps))): - # anneal temperature - temperature = start_temperature * (steps_until_x0 / max_steps) - # get noise level, per cosine scheduling - noise_p = math.cos( timestep * math.pi * 0.5 ) - # number of tokens to mask off to "noise" the input sequence - masked_tokens_n = max(int( noise_p * seq_len ), 1) - # pick the worst scoring tokens to mask off - masked_indices = scores.topk( masked_tokens_n, dim=-1 ).indices - # mask off inputs - input_ids = input_ids.scatter(0, masked_indices, self.stop_token) - # boolean mask - is_masked = input_ids == self.stop_token - # setup inputs - - inputs = _super.inputs( - text_list=[ text_list[batch_index] ] if text_list else None, - proms_list=[ proms_list[batch_index] ] if proms_list else None, - resps_list=[ input_ids ], - lang_list=[ lang_list[batch_index] ] if lang_list else None, - tone_list=[ tone_list[batch_index] ] if tone_list else None, - time_list=[ timestep ], - quant_levels=[ quant_levels[batch_index] ] if quant_levels else None, - ) - output = _super.forward( - inputs=inputs, - quant_levels=[ quant_levels[batch_index] ] if quant_levels else None, - #layer_skip_variables=sampling_layer_skip_variables, - ) - - logits = output.logits - - if cfg_strength > 0: - null_inputs = _super.inputs( - text_list=[ null_text ], - proms_list=[ null_prom ], - resps_list=[ input_ids ], - lang_list=[ lang_list[batch_index] ] if lang_list else None, - tone_list=[ tone_list[batch_index] ] if tone_list else None, - time_list=[ timestep ], - quant_levels=[ quant_levels[batch_index] ] if quant_levels else None, - ) - null_output = _super.forward( - inputs=null_inputs, - quant_levels=[ quant_levels[batch_index] ] if quant_levels else None, - #layer_skip_variables=sampling_layer_skip_variables, - ) - for logit, null_logit in zip(output.logits, null_output.logits): - logit[-seq_len:] = null_logit[-seq_len:] + ( logit[-seq_len:] - null_logit[-seq_len:] ) * cfg_strength - - # sample with sampler settings - filtered_sampled = _super.sample( - logits=logits, - prev_list=prev_list, - quant_levels=[ quant_levels[batch_index] ] if quant_levels else None, - - temperature=temperature, - min_temperature=sampling_min_temperature, - top_p=sampling_top_p, - top_k=sampling_top_k, - min_p=sampling_min_p, - repetition_penalty=sampling_repetition_penalty, - repetition_penalty_decay=sampling_repetition_penalty_decay, - length_penalty=sampling_length_penalty, - ) - - # retrieves unfiltered logits - unfiltered_sampled = _super.sample( - logits=logits, - prev_list=prev_list, - quant_levels=[ quant_levels[batch_index] ] if quant_levels else None, - temperature=0.0, - ) - # update previous list of tokens - prev_list = [ input_ids ] - - # extract logits - filtered_logits = filtered_sampled.logits[0] - unfiltered_logits = unfiltered_sampled.logits[0] - - # extract scores - filtered_scores = filtered_sampled.scores[0] - unfiltered_scores = unfiltered_sampled.scores[0] - - # extract sampled tokens - filtered_tokens = filtered_sampled[0][0] - unfiltered_tokens = unfiltered_sampled[0][0] - - # sample with gumbelnoise - # I actually feel like this doesn't matter? it's hard to judge with a partially trained NAR-len model - sampled_ids = gumbel_sample( filtered_logits, temperature=temperature, dim=-1 ) - #sampled_ids = filtered_tokens - - # keep unmasked tokens - input_ids = torch.where( is_masked, sampled_ids, input_ids ) - # update scores (conjugated to put the worst scores at the top) - scores = 1.0 - torch.tensor([score for score in unfiltered_scores], device=device) - - if cfg.experimental and max_steps > 0: - print( timestep, steps_until_x0, noise_p, masked_tokens_n, input_ids, scores ) - - return input_ids - - # perform demasked sampling (mock diffusion) - resps_list = [ demask_sampling( batch_index=i, seq_len=l ) for i, l in enumerate( len_list ) ] + resps_list = self.forward_nar_masked( + text_list=text_list, + proms_list=proms_list, + resps_list=resps_list, + task_list=task_list, + lang_list=lang_list, + tone_list=tone_list, + len_list=len_list, + **sampling_kwargs, + ) # expand if given a raw 1D tensor for i, resp in enumerate(resps_list): @@ -449,17 +425,7 @@ class AR_NAR(Base): logits=logits, prev_list=prev_list, quant_levels=quant_levels, - - temperature=sampling_temperature, - #min_temperature=sampling_min_temperature, - #top_p=sampling_top_p, - #top_k=sampling_top_k, - #min_p=sampling_min_p, - #repetition_penalty=sampling_repetition_penalty, - #repetition_penalty_decay=sampling_repetition_penalty_decay, - #length_penalty=sampling_length_penalty, - #beam_width=sampling_beam_width, - #mirostat=mirostat, + **sampling_kwargs, ) resps_list = sampled[0] @@ -478,41 +444,9 @@ class AR_NAR(Base): lang_list: list[Tensor] | None = None, tone_list: list[Tensor] | None = None, len_list: list[Tensor] | None = None, - - training: bool | int | None = None, - - max_steps: int = 1000, - max_levels: int = 0, - - input_prompt_prefix: bool = False, - prefix_silence: float = 1.0, - denoise_start: float = 0.0, - - sampling_temperature: float = 1.0, - sampling_min_temperature: float = -1.0, - sampling_top_k: int = -100, - sampling_top_p: float = 1.0, - sampling_min_p: float = 0.0, - sampling_repetition_penalty: float = 1.0, - sampling_repetition_penalty_decay: float = 0.0, - sampling_length_penalty: float = 0.0, - sampling_beam_width: int = 0, - sampling_mirostat_tau: float = 0.0, - sampling_mirostat_eta: float = 0.1, - sampling_dry_multiplier=0.0, - sampling_dry_base=1.75, - sampling_dry_allowed_length=2, - sampling_entropix=False, - - sampling_layer_skip: bool = False, - sampling_layer_skip_exit_layer: int = -1, - sampling_layer_skip_entropy_threshold: float = -1, - sampling_layer_skip_varentropy_threshold: float = -1, - - sampling_refine_on_stop: bool = False, - disable_tqdm=False, use_lora=None, + **sampling_kwargs, ): # deduce batch_size if text_list is not None: @@ -527,6 +461,21 @@ class AR_NAR(Base): if cfg.lora is not None: enable_lora( self, cfg.lora.active_level( 0 ) if use_lora is None else use_lora ) + # convert AR specific args + sampling_kwargs = convert_kwargs( sampling_kwargs, "ar_" ) + + temperature = sampling_kwargs.get("temperature", 1.0) + min_temperature = sampling_kwargs.get("min_temperature", -1.0) + max_duration = sampling_kwargs.get("max_duration", 500) + beam_width = sampling_kwargs.get("beam_width", 0) + entropix_sampling = sampling_kwargs.get("entropix_sampling", False) + refine_on_stop = sampling_kwargs.get("refine_on_stop", False) + input_prompt_prefix = sampling_kwargs.get("input_prompt_prefix", False) + layer_skip = sampling_kwargs.get("layer_skip", False) + prefix_silence = sampling_kwargs.get("prefix_silence", 0.0) + mirostat_tau = sampling_kwargs.get("mirostat_tau", 0.0) + mirostat_eta = sampling_kwargs.get("mirostat_eta", 0.0) + # inference len if task_list is not None and task_list[0] == "len": sequence_list = [ torch.tensor([0], device=device,dtype=torch.int16) for _ in range(batch_size) ] @@ -534,7 +483,7 @@ class AR_NAR(Base): stop_token = 10 task_list = [ "len" for _ in range(batch_size) ] - quant_levels = [ 0 for _ in range( max( batch_size, sampling_beam_width ) ) ] + quant_levels = [ 0 for _ in range( max( batch_size, beam_width ) ) ] for n in trange(10, desc="AR", disable=disable_tqdm): len_list = sequence_list @@ -586,22 +535,13 @@ class AR_NAR(Base): state = None mirostat = [ - {"n": 1024, "tau": sampling_mirostat_tau, "eta": sampling_mirostat_eta, "max_surprise": sampling_mirostat_eta * 2, "error_surprise": 0, "running_total_surprise": 0} - ] * batch_size if sampling_mirostat_tau > 0.0 else None + {"n": 1024, "tau": mirostat_tau, "eta": mirostat_eta, "max_surprise": mirostat_eta * 2, "error_surprise": 0, "running_total_surprise": 0} + ] * batch_size if mirostat_tau > 0.0 else None - scores = [ 1.0 ] * sampling_beam_width + scores = [ 1.0 ] * beam_width metrics = [] - # ick """ - low_temperature = False # sampling_temperature < 0.6 # sampling_repetition_penalty == 1.0 and sampling_temperature == 0.0 # - low_temperature_range = cfg.dataset.frames_per_second * 5 - - original_sampling_temperature = sampling_temperature - original_sampling_repetition_penalty = sampling_repetition_penalty - original_sampling_repetition_penalty_decay = sampling_repetition_penalty_decay - """ - sampling_layer_skip_variables = {} if sampling_layer_skip else None if sampling_layer_skip: @@ -611,6 +551,7 @@ class AR_NAR(Base): sampling_layer_skip_variables["varentropy_threshold"] = sampling_layer_skip_varentropy_threshold if sampling_layer_skip_exit_layer >= 0: sampling_layer_skip_variables["max_layer"] = sampling_layer_skip_exit_layer + """ for i, sequence in enumerate( sequence_list ): # add to text for STT @@ -627,23 +568,11 @@ class AR_NAR(Base): # start_slice[i] = sequence_list[i].shape[0] # get next in sequence - for n in trange(max_steps // max(1, self.causal_size), desc="AR", disable=disable_tqdm): + for n in trange(max_duration // max(1, self.causal_size), desc="AR", disable=disable_tqdm): # it would technically be faster to just append the new token's embedding to the inputs, but there's a VERY small performance gain from doing it, so it's not worth it text_list = [ sequence_list[i] if task in text_task else text_list[i] for i, task in enumerate(task_list) ] resps_list = [ sequence_list[i] if task not in text_task else resps_list[i] for i, task in enumerate(task_list) ] - # greedy sampling in the AR *does* work, but requires some quasi-exotic sampling to work around the initial burst of garbage from polluting the rest of the sequence - # naturally, rep pen wrangles this initial burst of noise, but naively relying on rep_pen is no good, as it fails after ~6 seconds of audio - # however, switching to a default sampling temperature with "clean greedy sampled codes" will make the rest of sequence sound as if it were greedy sampled - # to-do: tune these values, maybe have it factor based on confidence scores or something - """ - if low_temperature: - enabled = n < low_temperature_range - sampling_repetition_penalty = 1.125 if enabled else 1.25 - #sampling_repetition_penalty_decay = 0.0 if enabled else original_sampling_repetition_penalty_decay - #sampling_temperature = original_sampling_temperature if enabled else 1.0 - """ - inputs = self.inputs( text_list=text_list, proms_list=proms_list, @@ -652,7 +581,7 @@ class AR_NAR(Base): tone_list=tone_list, len_list=len_list, task_list=task_list, - quant_levels=[ 0 for _ in range( max( batch_size, sampling_beam_width ) ) ] + quant_levels=[ 0 for _ in range( max( batch_size, beam_width ) ) ] ) # to-do: find an elegant way to write this @@ -660,31 +589,14 @@ class AR_NAR(Base): inputs=inputs, state=state, #layer_skip_variables=sampling_layer_skip_variables, - output_attentions=sampling_entropix, + output_attentions=entropix_sampling, ) logits, state = output.logits, output.state sampled = super().sample( logits=logits, - prev_list=None if sampling_repetition_penalty == 1.0 and sampling_length_penalty == 0.0 else [ resps_list[i] if task not in text_task else text_list[i] for i, task in enumerate( task_list ) ], - - temperature=sampling_temperature, - min_temperature=sampling_min_temperature, - top_p=sampling_top_p, - top_k=sampling_top_k, - min_p=sampling_min_p, - repetition_penalty=sampling_repetition_penalty, - repetition_penalty_decay=sampling_repetition_penalty_decay, - length_penalty=sampling_length_penalty, - beam_width=sampling_beam_width, - - mirostat=mirostat, - - dry_multiplier=sampling_dry_multiplier, - dry_base=sampling_dry_base, - dry_allowed_length=sampling_dry_allowed_length, - - attentions=output.attentions if sampling_entropix else None, + prev_list=[ resps_list[i] if task not in text_task else text_list[i] for i, task in enumerate( task_list ) ], + **(sampling_kwargs | {"attentions": output.attentions if entropix_sampling else None}), ) r = sampled[0] @@ -698,17 +610,17 @@ class AR_NAR(Base): if mirostat is not None: mirostat = sampled.scores - elif sampling_beam_width > 0: + elif beam_width > 0: # expand tuple s = sampled.scores # first step, expand batch if batch_size == 1: - batch_size = sampling_beam_width - text_list = text_list * sampling_beam_width - proms_list = proms_list * sampling_beam_width - sequence_list = sequence_list * sampling_beam_width - task_list = task_list * sampling_beam_width - start_slice = start_slice * sampling_beam_width + batch_size = beam_width + text_list = text_list * beam_width + proms_list = proms_list * beam_width + sequence_list = sequence_list * beam_width + task_list = task_list * beam_width + start_slice = start_slice * beam_width stopped = torch.zeros(batch_size, device=device).bool() scores = [ scores[i] + score for i, score in enumerate(s) ] @@ -727,22 +639,21 @@ class AR_NAR(Base): break # to-do for layerskip / speculative sampling: rerun the last sequence again at max depth - + """ if metrics: from ..plot import plot_sample_metrics filename = "metrics" - if sampling_entropix: - filename += f'[entropix]' - """ + if entropix_sampling: + filename += f'[entropix_sampling]' if sampling_layer_skip_exit_layer >= 0: filename += f'[{sampling_layer_skip_exit_layer+1}]' - """ plot_sample_metrics( metrics, filename=f'{filename}.png' ) + """ # pick the best scoring candidate # desu this is always going to be candidate 0 - if sampling_beam_width: + if beam_width: sequence_list = sequence_list[:1] task_list = task_list[:1] @@ -751,7 +662,7 @@ class AR_NAR(Base): # remove sequence_list = [ sequence_list[i][start_slice[i]:] for i, task in enumerate( task_list ) ] - if sampling_refine_on_stop: + if refine_on_stop: # get how much we need to slice from the end slice_lengths = [ sequence.shape[-1] for sequence in sequence_list ] # -1 for the stop token @@ -777,69 +688,10 @@ class AR_NAR(Base): training: bool | int | None = None, - max_steps: int = 1000, - max_levels: int = 0, - - input_prompt_prefix: bool = False, - prefix_silence: float = 1.0, - denoise_start: float = 0.0, - - sampling_temperature: float = 1.0, - sampling_min_temperature: float = -1.0, - sampling_top_k: int = -100, - sampling_top_p: float = 1.0, - sampling_min_p: float = 0.0, - sampling_repetition_penalty: float = 1.0, - sampling_repetition_penalty_decay: float = 0.0, - sampling_length_penalty: float = 0.0, - sampling_beam_width: int = 0, - sampling_mirostat_tau: float = 0.0, - sampling_mirostat_eta: float = 0.1, - sampling_dry_multiplier=0.0, - sampling_dry_base=1.75, - sampling_dry_allowed_length=2, - sampling_entropix=False, - - sampling_layer_skip: bool = False, - sampling_layer_skip_exit_layer: int = -1, - sampling_layer_skip_entropy_threshold: float = -1, - sampling_layer_skip_varentropy_threshold: float = -1, - - sampling_refine_on_stop: bool = False, - disable_tqdm=False, use_lora=None, + **sampling_kwargs, ): - kwargs = dict( - max_steps=max_steps, - max_levels=max_levels, - input_prompt_prefix=input_prompt_prefix, - prefix_silence=prefix_silence, - denoise_start=denoise_start, - sampling_temperature=sampling_temperature, - sampling_min_temperature=sampling_min_temperature, - sampling_top_k=sampling_top_k, - sampling_top_p=sampling_top_p, - sampling_min_p=sampling_min_p, - sampling_repetition_penalty=sampling_repetition_penalty, - sampling_repetition_penalty_decay=sampling_repetition_penalty_decay, - sampling_length_penalty=sampling_length_penalty, - sampling_beam_width=sampling_beam_width, - sampling_mirostat_tau=sampling_mirostat_tau, - sampling_mirostat_eta=sampling_mirostat_eta, - sampling_dry_multiplier=sampling_dry_multiplier, - sampling_dry_base=sampling_dry_base, - sampling_dry_allowed_length=sampling_dry_allowed_length, - sampling_entropix=sampling_entropix, - sampling_layer_skip=sampling_layer_skip, - sampling_layer_skip_exit_layer=sampling_layer_skip_exit_layer, - sampling_layer_skip_entropy_threshold=sampling_layer_skip_entropy_threshold, - sampling_layer_skip_varentropy_threshold=sampling_layer_skip_varentropy_threshold, - sampling_refine_on_stop=sampling_refine_on_stop, - disable_tqdm=disable_tqdm, - use_lora=use_lora, - ) - # deduce batch_size if text_list is not None: default_task = "tts" @@ -883,7 +735,7 @@ class AR_NAR(Base): lang_list=lang_list, tone_list=tone_list, len_list=len_list, - **kwargs, + **sampling_kwargs, ) # is AR @@ -895,7 +747,7 @@ class AR_NAR(Base): lang_list=lang_list, tone_list=tone_list, len_list=len_list, - **kwargs, + **sampling_kwargs, ) @@ -1081,12 +933,12 @@ def example_usage(): text_list, proms_list, resp_list, task_list = sample_data( task ) if task == "tts-nar": - len_list = engine(text_list, proms_list, task_list=["len"], max_steps=5, sampling_temperature=0.0 ) + len_list = engine(text_list, proms_list, task_list=["len"], max_steps=5, temperature=0.0 ) len_list = [ resp_list[0].shape[0] for l in len_list ] - resps_list = engine( text_list, proms_list, len_list=len_list, sampling_temperature=0.0 ) + resps_list = engine( text_list, proms_list, len_list=len_list, temperature=0.0 ) else: - resps_list = engine( text_list, proms_list, task_list=["tts"], max_steps=steps, sampling_temperature=1.0 ) - resps_list = engine( text_list, proms_list, resps_list=resps_list, sampling_temperature=0.0 ) + resps_list = engine( text_list, proms_list, task_list=["tts"], max_duration=steps, temperature=1.0 ) + resps_list = engine( text_list, proms_list, resps_list=resps_list, temperature=0.0 ) for i, o in enumerate(resps_list): _ = decode_to_file(o.to(dtype=torch.int32), f"data/{cfg.model.arch_type}.{cfg.audio_backend}.{i}.{name}.{task}.wav", device=cfg.device) diff --git a/vall_e/models/base.py b/vall_e/models/base.py index a68eae7..0867e3d 100755 --- a/vall_e/models/base.py +++ b/vall_e/models/base.py @@ -1680,29 +1680,31 @@ class Base(nn.Module): self, logits: list[Tensor], # logit scores prev_list: list[Tensor] | None = None, # previous tokens - quant_levels: int | list[int] | Tensor | None = None, - # base sampling parameters - temperature: float = 1.0, - min_temperature: float = -1.0, # activates dynamic temperature sampling - top_k: int = -100, - top_p: float = 1.0, - min_p: float = 0.0, - # repetition penalty parameters - repetition_penalty: float = 1.0, - repetition_penalty_decay: float = 0.0, - # length penalty parameters - length_penalty: float = 0.0, - # beam sampling parameters - beam_width: int = 0, - # mirostat sampling parameters - mirostat: list[dict] | None = None, - # DRY sampling parameters - dry_multiplier=0.0, - dry_base=1.75, - dry_allowed_length=2, - # other - attentions=None, + quant_levels: int | list[int] | Tensor | None = None, + **sampling_kwargs, ): + # yikes + temperature = sampling_kwargs.get("temperature", 1.0) + min_temperature = sampling_kwargs.get("min_temperature", -1.0) + top_k = sampling_kwargs.get("top_k", -100) + top_p = sampling_kwargs.get("top_p", 1.0) + min_p = sampling_kwargs.get("min_p", 0.0) + # repetition penalty parameters + repetition_penalty = sampling_kwargs.get("repetition_penalty", 1.0) + repetition_penalty_decay = sampling_kwargs.get("repetition_penalty_decay", 0.0) + # length penalty parameters + length_penalty = sampling_kwargs.get("length_penalty", 0.0) + # beam sampling parameters + beam_width = sampling_kwargs.get("beam_width", 0) + # mirostat sampling parameters + mirostat = sampling_kwargs.get("mirostat", None) + # DRY sampling parameters + dry_multiplier = sampling_kwargs.get("dry_multiplier", 0.0) + dry_base = sampling_kwargs.get("dry_base", 1.75) + dry_allowed_length = sampling_kwargs.get("dry_allowed_length", 2) + + attentions = sampling_kwargs.get("attentions", None) + batch_size = len( logits ) if min_temperature < 0: diff --git a/vall_e/utils/__init__.py b/vall_e/utils/__init__.py index 367e35a..3dbe232 100755 --- a/vall_e/utils/__init__.py +++ b/vall_e/utils/__init__.py @@ -14,5 +14,6 @@ from .utils import ( timer, prune_missing, clamp, - md5_hash + md5_hash, + convert_kwargs ) \ No newline at end of file diff --git a/vall_e/utils/utils.py b/vall_e/utils/utils.py index 9f85087..6e789ef 100755 --- a/vall_e/utils/utils.py +++ b/vall_e/utils/utils.py @@ -32,11 +32,27 @@ from datetime import datetime T = TypeVar("T") +# removes prefix from key in a dict +# useful for mapping args like ar_temperature => temperature +def convert_kwargs( kwargs, prefix ): + copied = {} | kwargs + + for key, value in copied.items(): + if not key.startswith( prefix ): + continue + + kwargs.pop(key) + kwargs[key[len(prefix):]] = value + + return kwargs + +# hashes values or a list of values def md5_hash( x ): if isinstance( x, list ): return md5_hash(":".join([ md5_hash( _ ) for _ in x ])) return hashlib.md5(str(x).encode("utf-8")).hexdigest() +# removes entries from a dict if that key is missing from the source def prune_missing( source, dest, recurse=True, path=[], parent_is_obj=None, return_missing=True ): is_obj = hasattr( source, "__dict__" ) if parent_is_obj is None: diff --git a/vall_e/webui.py b/vall_e/webui.py index 9e16532..c8bbbaf 100644 --- a/vall_e/webui.py +++ b/vall_e/webui.py @@ -192,11 +192,11 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ): raise Exception("No model loaded.") if kwargs.pop("dynamic-sampling", False): - kwargs['min-ar-temp'] = 0.01 if kwargs['ar-temp'] > 0.01 else 0.0 - kwargs['min-nar-temp'] = 0.0 # 0.85 if kwargs['nar-temp'] > 0.85 else 0.0 # should probably disable it for the NAR + kwargs['min-ar-temperature'] = 0.01 if kwargs['ar-temperature'] > 0.01 else 0.0 + kwargs['min-nar-temperature'] = 0.0 # 0.85 if kwargs['nar-temperature'] > 0.85 else 0.0 # should probably disable it for the NAR else: - kwargs['min-ar-temp'] = -1 - kwargs['min-nar-temp'] = -1 + kwargs['min-ar-temperature'] = -1 + kwargs['min-nar-temperature'] = -1 parser = argparse.ArgumentParser(allow_abbrev=False, add_help=False) # I'm very sure I can procedurally generate this list @@ -205,14 +205,15 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ): parser.add_argument("--references", type=str, default=kwargs["reference"]) parser.add_argument("--language", type=str, default=kwargs["language"]) parser.add_argument("--input-prompt-length", type=float, default=kwargs["input-prompt-length"]) - parser.add_argument("--input-prompt-prefix", action='store_true', default=kwargs["input-prompt-prefix"] if cfg.experimental else False) - parser.add_argument("--max-ar-steps", type=int, default=int(kwargs["max-seconds"]*cfg.dataset.frames_per_second)) - parser.add_argument("--max-nar-levels", type=int, default=kwargs["max-nar-levels"] if cfg.experimental else 0) - parser.add_argument("--ar-temp", type=float, default=kwargs["ar-temp"]) - parser.add_argument("--nar-temp", type=float, default=kwargs["nar-temp"]) - parser.add_argument("--min-ar-temp", type=float, default=kwargs["min-ar-temp"]) - parser.add_argument("--min-nar-temp", type=float, default=kwargs["min-nar-temp"]) - parser.add_argument("--prefix-silence", type=float, default=kwargs["prefix-silence"] if cfg.experimental else 0) + parser.add_argument("--input-prompt-prefix", action='store_true', default=kwargs["input-prompt-prefix"]) + parser.add_argument("--max-duration", type=int, default=int(kwargs["max-duration"]*cfg.dataset.frames_per_second)) + parser.add_argument("--max-levels", type=int, default=kwargs["max-levels"]) + parser.add_argument("--max-steps", type=int, default=kwargs["max-steps"]) + parser.add_argument("--ar-temperature", type=float, default=kwargs["ar-temperature"]) + parser.add_argument("--nar-temperature", type=float, default=kwargs["nar-temperature"]) + parser.add_argument("--min-ar-temperature", type=float, default=kwargs["min-ar-temperature"]) + parser.add_argument("--min-nar-temperature", type=float, default=kwargs["min-nar-temperature"]) + parser.add_argument("--prefix-silence", type=float, default=kwargs["prefix-silence"]) parser.add_argument("--top-p", type=float, default=kwargs["top-p"]) parser.add_argument("--top-k", type=int, default=kwargs["top-k"]) parser.add_argument("--min-p", type=float, default=kwargs["min-p"]) @@ -227,10 +228,11 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ): parser.add_argument("--dry-allowed-length", type=int, default=kwargs["dry-allowed-length"]) parser.add_argument("--entropix-sampling", action="store_true") parser.add_argument("--layer-skip", action="store_true") - parser.add_argument("--layer-skip-exit-layer", type=int, default=kwargs["layer-skip-exit-layer"] if cfg.experimental else -1) - parser.add_argument("--layer-skip-entropy-threshold", type=int, default=kwargs["layer-skip-entropy-threshold"] if cfg.experimental else 0.1) - parser.add_argument("--layer-skip-varentropy-threshold", type=int, default=kwargs["layer-skip-varentropy-threshold"] if cfg.experimental else 0.1) + parser.add_argument("--layer-skip-exit-layer", type=int, default=kwargs["layer-skip-exit-layer"]) + parser.add_argument("--layer-skip-entropy-threshold", type=int, default=kwargs["layer-skip-entropy-threshold"]) + parser.add_argument("--layer-skip-varentropy-threshold", type=int, default=kwargs["layer-skip-varentropy-threshold"]) parser.add_argument("--refine-on-stop", action="store_true") + parser.add_argument("--denoise-start", type=float, default=0.0) args, unknown = parser.parse_known_args() if is_windows: @@ -255,6 +257,27 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ): tts = init_tts() gr.Info("Inferencing...") + + sampling_kwargs = dict( + max_duration=args.max_duration, + ar_temperature=args.ar_temperature, nar_temperature=args.nar_temperature, + min_ar_temperature=args.min_ar_temperature, min_nar_temperature=args.min_nar_temperature, + top_p=args.top_p, top_k=args.top_k, min_p=args.min_p, + repetition_penalty=args.repetition_penalty, repetition_penalty_decay=args.repetition_penalty_decay, + length_penalty=args.length_penalty, + beam_width=args.beam_width, + mirostat_tau=args.mirostat_tau, mirostat_eta=args.mirostat_eta, + dry_multiplier=args.dry_multiplier, dry_base=args.dry_base, dry_allowed_length=args.dry_allowed_length, + entropix_sampling=args.entropix_sampling, + layer_skip=args.layer_skip, + layer_skip_exit_layer=args.layer_skip_exit_layer, + layer_skip_entropy_threshold=args.layer_skip_entropy_threshold, + layer_skip_varentropy_threshold=args.layer_skip_varentropy_threshold, + refine_on_stop=args.refine_on_stop, + denoise_start=args.denoise_start, + prefix_silence=args.prefix_silence, + input_prompt_prefix=args.input_prompt_prefix, + ) with timer("Inferenced in", callback=lambda msg: gr.Info( msg )) as t: wav, sr = tts.inference( @@ -262,34 +285,7 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ): language=args.language, task=args.task, references=args.references.split(";") if args.references is not None else [], - out_path=tmp.name, - max_ar_steps=args.max_ar_steps, - max_nar_levels=args.max_nar_levels, - input_prompt_length=args.input_prompt_length, - input_prompt_prefix=args.input_prompt_prefix, - prefix_silence=args.prefix_silence, - ar_temp=args.ar_temp, - nar_temp=args.nar_temp, - min_ar_temp=args.min_ar_temp, - min_nar_temp=args.min_nar_temp, - top_p=args.top_p, - top_k=args.top_k, - min_p=args.min_p, - beam_width=args.beam_width, - repetition_penalty=args.repetition_penalty, - repetition_penalty_decay=args.repetition_penalty_decay, - length_penalty=args.length_penalty, - mirostat_tau=args.mirostat_tau, - mirostat_eta=args.mirostat_eta, - dry_multiplier=args.dry_multiplier, - dry_base=args.dry_base, - dry_allowed_length=args.dry_allowed_length, - entropix_sampling=args.entropix_sampling, - - layer_skip=args.layer_skip, - layer_skip_entropy_threshold=args.layer_skip_entropy_threshold, - layer_skip_varentropy_threshold=args.layer_skip_varentropy_threshold, - refine_on_stop=args.refine_on_stop, + **sampling_kwargs, ) wav = wav.squeeze(0).cpu().numpy() @@ -301,20 +297,28 @@ def do_inference_stt( progress=gr.Progress(track_tqdm=True), *args, **kwargs ): raise Exception("No model loaded.") if kwargs.pop("dynamic-sampling", False): - kwargs['min-ar-temp'] = 0.85 if kwargs['ar-temp'] > 0.85 else 0.0 + kwargs['min-ar-temperature'] = 0.85 if kwargs['ar-temperature'] > 0.85 else 0.0 else: - kwargs['min-ar-temp'] = -1 + kwargs['min-ar-temperature'] = -1 parser = argparse.ArgumentParser(allow_abbrev=False, add_help=False) # I'm very sure I can procedurally generate this list + parser.add_argument("--text", type=str, default=kwargs["text"]) + parser.add_argument("--task", type=str, default="tts") parser.add_argument("--references", type=str, default=kwargs["reference"]) parser.add_argument("--language", type=str, default=kwargs["language"]) - parser.add_argument("--max-ar-steps", type=int, default=0) - parser.add_argument("--ar-temp", type=float, default=kwargs["ar-temp"]) - parser.add_argument("--min-ar-temp", type=float, default=kwargs["min-ar-temp"]) + parser.add_argument("--input-prompt-length", type=float, default=kwargs["input-prompt-length"]) + parser.add_argument("--input-prompt-prefix", action='store_true', default=kwargs["input-prompt-prefix"]) + parser.add_argument("--max-duration", type=int, default=int(kwargs["max-duration"]*cfg.dataset.frames_per_second)) + parser.add_argument("--max-levels", type=int, default=kwargs["max-levels"]) + parser.add_argument("--ar-temperature", type=float, default=kwargs["ar-temperature"]) + parser.add_argument("--nar-temperature", type=float, default=kwargs["nar-temperature"]) + parser.add_argument("--min-ar-temperature", type=float, default=kwargs["min-ar-temperature"]) + parser.add_argument("--min-nar-temperature", type=float, default=kwargs["min-nar-temperature"]) + parser.add_argument("--prefix-silence", type=float, default=kwargs["prefix-silence"]) parser.add_argument("--top-p", type=float, default=kwargs["top-p"]) parser.add_argument("--top-k", type=int, default=kwargs["top-k"]) - parser.add_argument("--min-p", type=int, default=kwargs["min-p"]) + parser.add_argument("--min-p", type=float, default=kwargs["min-p"]) parser.add_argument("--repetition-penalty", type=float, default=kwargs["repetition-penalty"]) parser.add_argument("--repetition-penalty-decay", type=float, default=kwargs["repetition-penalty-decay"]) parser.add_argument("--length-penalty", type=float, default=kwargs["length-penalty"]) @@ -325,6 +329,12 @@ def do_inference_stt( progress=gr.Progress(track_tqdm=True), *args, **kwargs ): parser.add_argument("--dry-base", type=float, default=kwargs["dry-base"]) parser.add_argument("--dry-allowed-length", type=int, default=kwargs["dry-allowed-length"]) parser.add_argument("--entropix-sampling", action="store_true") + parser.add_argument("--layer-skip", action="store_true") + parser.add_argument("--layer-skip-exit-layer", type=int, default=kwargs["layer-skip-exit-layer"]) + parser.add_argument("--layer-skip-entropy-threshold", type=int, default=kwargs["layer-skip-entropy-threshold"]) + parser.add_argument("--layer-skip-varentropy-threshold", type=int, default=kwargs["layer-skip-varentropy-threshold"]) + parser.add_argument("--refine-on-stop", action="store_true") + parser.add_argument("--cfg-strength", type=float, default=kwargs["cfg-strength"]) args, unknown = parser.parse_known_args() @@ -334,17 +344,36 @@ def do_inference_stt( progress=gr.Progress(track_tqdm=True), *args, **kwargs ): """ args.references = args.references.split(";") if args.references is not None else [] - if args.max_ar_steps == 0: + if args.max_duration == 0: for i, path in enumerate( args.references ): metadata = torchaudio.info(path) duration = metadata.num_frames / metadata.sample_rate - args.max_ar_steps += duration - args.max_ar_steps = math.floor( args.max_ar_steps * 20 ) # assume 20 tokens per second + args.max_duration += duration + args.max_duration = math.floor( args.max_duration * 20 ) # assume 20 tokens per second if kwargs.pop("entropix-sampling", False): args.entropix_sampling = True tts = init_tts() + + sampling_kwargs = dict( + max_duration=args.max_duration, + ar_temperature=args.ar_temperature, nar_temperature=args.nar_temperature, + min_ar_temperature=args.min_ar_temperature, min_nar_temperature=args.min_nar_temperature, + top_p=args.top_p, top_k=args.top_k, min_p=args.min_p, + repetition_penalty=args.repetition_penalty, repetition_penalty_decay=args.repetition_penalty_decay, + length_penalty=args.length_penalty, + beam_width=args.beam_width, + mirostat_tau=args.mirostat_tau, mirostat_eta=args.mirostat_eta, + dry_multiplier=args.dry_multiplier, dry_base=args.dry_base, dry_allowed_length=args.dry_allowed_length, + entropix_sampling=args.entropix_sampling, + layer_skip=args.layer_skip, + layer_skip_exit_layer=args.layer_skip_exit_layer, + layer_skip_entropy_threshold=args.layer_skip_entropy_threshold, + layer_skip_varentropy_threshold=args.layer_skip_varentropy_threshold, + refine_on_stop=args.refine_on_stop, + denoise_start=args.denoise_start, + ) gr.Info("Inferencing...") with timer("Inferenced in") as t: @@ -353,21 +382,7 @@ def do_inference_stt( progress=gr.Progress(track_tqdm=True), *args, **kwargs ): language=args.language, task="stt", references=args.references, - max_ar_steps=args.max_ar_steps, - ar_temp=args.ar_temp, - min_ar_temp=args.min_ar_temp, - top_p=args.top_p, - top_k=args.top_k, - min_p=args.min_p, - repetition_penalty=args.repetition_penalty, - repetition_penalty_decay=args.repetition_penalty_decay, - length_penalty=args.length_penalty, - mirostat_tau=args.mirostat_tau, - mirostat_eta=args.mirostat_eta, - dry_multiplier=args.dry_multiplier, - dry_base=args.dry_base, - dry_allowed_length=args.dry_allowed_length, - entropix_sampling=args.entropix_sampling, + **sampling_kwargs, ) return text @@ -424,12 +439,13 @@ with ui: with gr.Column(scale=7): with gr.Tab("Basic Settings"): with gr.Row(): - layout["inference_tts"]["inputs"]["max-seconds"] = gr.Slider(value=12, minimum=1, maximum=32, step=0.1, label="Maximum Seconds", info="Limits how many steps to perform in the AR pass.") + layout["inference_tts"]["inputs"]["max-duration"] = gr.Slider(value=12, minimum=1, maximum=32, step=0.1, label="Maximum Seconds", info="Limits how many steps to perform in the AR pass.") layout["inference_tts"]["inputs"]["input-prompt-length"] = gr.Slider(value=5.0, minimum=0.0, maximum=12.0, step=0.05, label="Input Prompt Repeat/Trim Length", info="Repeats and trims the input prompt down to X seconds. Set 0 to disable.") with gr.Row(): - layout["inference_tts"]["inputs"]["ar-temp"] = gr.Slider(value=0.5, minimum=0.0, maximum=1.5, step=0.05, label="Temperature (AR)", info="Modifies the randomness from the samples in the AR. (0 to greedy* sample)") - layout["inference_tts"]["inputs"]["nar-temp"] = gr.Slider(value=0.0, minimum=0.0, maximum=1.5, step=0.05, label="Temperature (NAR)", info="Modifies the randomness from the samples in the NAR. (0 to greedy sample)") + layout["inference_tts"]["inputs"]["ar-temperature"] = gr.Slider(value=1.0, minimum=0.0, maximum=1.5, step=0.05, label="Temperature (AR)", info="Modifies the randomness from the samples in the AR. (0 to greedy* sample)") + layout["inference_tts"]["inputs"]["nar-temperature"] = gr.Slider(value=0.0, minimum=0.0, maximum=1.5, step=0.05, label="Temperature (NAR)", info="Modifies the randomness from the samples in the NAR. (0 to greedy sample)") with gr.Row(): + layout["inference_tts"]["inputs"]["cfg-strength"] = gr.Slider(value=0.0, minimum=0.0, maximum=3.0, step=0.05, label="CFG Strength", info="Classifier Free Guidance scale") layout["inference_tts"]["inputs"]["language"] = gr.Dropdown(choices=get_languages(), label="Language", value="en") with gr.Tab("Sampler Settings"): with gr.Row(): @@ -438,7 +454,7 @@ with ui: layout["inference_tts"]["inputs"]["min-p"] = gr.Slider(value=0.0, minimum=0.0, maximum=1.0, step=0.05, label="Min P") layout["inference_tts"]["inputs"]["beam-width"] = gr.Slider(value=0, minimum=0, maximum=32, step=1, label="Beam Width", info="Number of branches to search through for beam search sampling.") with gr.Row(): - layout["inference_tts"]["inputs"]["repetition-penalty"] = gr.Slider(value=1.5, minimum=-2.0, maximum=2.0, step=0.05, label="Repetition Penalty", info="Incurs a penalty to tokens based on how often they appear in a sequence.") + layout["inference_tts"]["inputs"]["repetition-penalty"] = gr.Slider(value=1.0, minimum=-2.0, maximum=2.0, step=0.05, label="Repetition Penalty", info="Incurs a penalty to tokens based on how often they appear in a sequence.") layout["inference_tts"]["inputs"]["repetition-penalty-decay"] = gr.Slider(value=0.0, minimum=-2.0, maximum=2.0, step=0.05, label="Repetition Penalty Length Decay", info="Modifies the reptition penalty based on how far back in time the token appeared in the sequence.") layout["inference_tts"]["inputs"]["length-penalty"] = gr.Slider(value=0.0, minimum=-2.0, maximum=2.0, step=0.05, label="Length Penalty", info="(AR only) Modifies the probability of a stop token based on the current length of the sequence.") with gr.Row(): @@ -448,24 +464,24 @@ with ui: layout["inference_tts"]["inputs"]["dry-multiplier"] = gr.Slider(value=0.0, minimum=0.0, maximum=8.0, step=0.05, label="DRY Multiplier", info="The multiplying factor for the DRY score penalty (0 to disable DRY sampling).") layout["inference_tts"]["inputs"]["dry-base"] = gr.Slider(value=1.75, minimum=0.0, maximum=8.0, step=0.05, label="DRY Base", info="The base of the exponent in the DRY score penalty") layout["inference_tts"]["inputs"]["dry-allowed-length"] = gr.Slider(value=2, minimum=0, maximum=75, step=1, label="Allowed Length", info="The maximimum length a token can be to perform DRY penalty with.") - if cfg.experimental: - with gr.Tab("Experimental Settings"): - with gr.Row(): - layout["inference_tts"]["inputs"]["max-nar-levels"] = gr.Slider(value=7, minimum=0, maximum=7, step=1, label="Max NAR Levels", info="Limits how many steps to perform in the NAR pass.") - layout["inference_tts"]["inputs"]["input-prompt-prefix"] = gr.Checkbox(label="Input Prompt as Prefix", info="Treats the input prompt clip as the prefix of the generated sequence.") - with gr.Row(): - layout["inference_tts"]["inputs"]["prefix-silence"] = gr.Slider(value=0.0, minimum=0.0, maximum=1.0, step=0.05, label="Silence Prefix Duration", info="Amount of silence to prefix to the output response before beginning inference.") - with gr.Row(): - layout["inference_tts"]["inputs"]["dynamic-sampling"] = gr.Checkbox(label="Dynamic Temperature", info="Dynamically adjusts the temperature based on the highest confident predicted token per sampling step.") - layout["inference_tts"]["inputs"]["entropix-sampling"] = gr.Checkbox(label="Entropix Sampling", info="Dynamically samples based on entropy/varentropy values from the logits / attention scores.") - with gr.Row(): - layout["inference_tts"]["inputs"]["layer-skip"] = gr.Checkbox(label="Layer Skip", info="Performs self-speculative early exit 'sampling'") - layout["inference_tts"]["inputs"]["refine-on-stop"] = gr.Checkbox(label="Refine on ", info="Uses the last step's logits for the AR sequence instead.") - with gr.Row(): - layout["inference_tts"]["inputs"]["layer-skip-exit-layer"] = gr.Slider(value=11, minimum=0, maximum=11, step=1, label="Layer Skip Exit Layer", info="Maximum model layer to exit early from.") - layout["inference_tts"]["inputs"]["layer-skip-entropy-threshold"] = gr.Slider(value=0.1, minimum=0, maximum=1.0, step=0.01, label="Layer Skip Entropy Threshold", info="Entropy threshold for early-exit") - layout["inference_tts"]["inputs"]["layer-skip-varentropy-threshold"] = gr.Slider(value=0.1, minimum=0, maximum=1.0, step=0.01, label="Layer Skip Varentropy Threshold", info="Varentropy threshold for early-exit") - + with gr.Tab("Experimental Settings", visible=cfg.experimental): + with gr.Row(): + layout["inference_tts"]["inputs"]["max-steps"] = gr.Slider(value=25, minimum=1, maximum=50, step=1, label="Max NAR Steps", info="Limits how many steps to perform in the NAR (demask) pass.") + layout["inference_tts"]["inputs"]["max-levels"] = gr.Slider(value=7, minimum=0, maximum=7, step=1, label="Max NAR Levels", info="Limits how many steps to perform in the NAR pass.") + layout["inference_tts"]["inputs"]["input-prompt-prefix"] = gr.Checkbox(label="Input Prompt as Prefix", info="Treats the input prompt clip as the prefix of the generated sequence.") + with gr.Row(): + layout["inference_tts"]["inputs"]["prefix-silence"] = gr.Slider(value=0.0, minimum=0.0, maximum=1.0, step=0.05, label="Silence Prefix Duration", info="Amount of silence to prefix to the output response before beginning inference.") + with gr.Row(): + layout["inference_tts"]["inputs"]["dynamic-sampling"] = gr.Checkbox(label="Dynamic Temperature", info="Dynamically adjusts the temperature based on the highest confident predicted token per sampling step.") + layout["inference_tts"]["inputs"]["entropix-sampling"] = gr.Checkbox(label="Entropix Sampling", info="Dynamically samples based on entropy/varentropy values from the logits / attention scores.") + with gr.Row(): + layout["inference_tts"]["inputs"]["layer-skip"] = gr.Checkbox(label="Layer Skip", info="Performs self-speculative early exit 'sampling'") + layout["inference_tts"]["inputs"]["refine-on-stop"] = gr.Checkbox(label="Refine on ", info="Uses the last step's logits for the AR sequence instead.") + with gr.Row(): + layout["inference_tts"]["inputs"]["layer-skip-exit-layer"] = gr.Slider(value=11, minimum=0, maximum=11, step=1, label="Layer Skip Exit Layer", info="Maximum model layer to exit early from.") + layout["inference_tts"]["inputs"]["layer-skip-entropy-threshold"] = gr.Slider(value=0.1, minimum=0, maximum=1.0, step=0.01, label="Layer Skip Entropy Threshold", info="Entropy threshold for early-exit") + layout["inference_tts"]["inputs"]["layer-skip-varentropy-threshold"] = gr.Slider(value=0.1, minimum=0, maximum=1.0, step=0.01, label="Layer Skip Varentropy Threshold", info="Varentropy threshold for early-exit") + layout["inference_tts"]["buttons"]["inference"].click( fn=do_inference_tts, @@ -485,7 +501,7 @@ with ui: with gr.Column(scale=7): with gr.Tab("Basic Settings"): with gr.Row(): - layout["inference_stt"]["inputs"]["ar-temp"] = gr.Slider(value=0.0, minimum=0.0, maximum=1.5, step=0.05, label="Temperature (AR)", info="Modifies the randomness from the samples in the AR. (0 to greedy sample)") + layout["inference_stt"]["inputs"]["ar-temperature"] = gr.Slider(value=0.0, minimum=0.0, maximum=1.5, step=0.05, label="Temperature (AR)", info="Modifies the randomness from the samples in the AR. (0 to greedy sample)") with gr.Row(): layout["inference_stt"]["inputs"]["dynamic-sampling"] = gr.Checkbox(label="Dynamic Temperature", info="Dynamically adjusts the temperature based on the highest confident predicted token per sampling step.") layout["inference_stt"]["inputs"]["language"] = gr.Dropdown(choices=get_languages(), label="Language", value="en") @@ -496,7 +512,7 @@ with ui: layout["inference_stt"]["inputs"]["min-p"] = gr.Slider(value=0.0, minimum=0.0, maximum=1.0, step=0.05, label="Min P") layout["inference_stt"]["inputs"]["beam-width"] = gr.Slider(value=0, minimum=0, maximum=32, step=1, label="Beam Width", info="Number of branches to search through for beam search sampling.") with gr.Row(): - layout["inference_stt"]["inputs"]["repetition-penalty"] = gr.Slider(value=1.25, minimum=-2.0, maximum=2.0, step=0.05, label="Repetition Penalty", info="Incurs a penalty to tokens based on how often they appear in a sequence.") + layout["inference_stt"]["inputs"]["repetition-penalty"] = gr.Slider(value=1.0, minimum=-2.0, maximum=2.0, step=0.05, label="Repetition Penalty", info="Incurs a penalty to tokens based on how often they appear in a sequence.") layout["inference_stt"]["inputs"]["repetition-penalty-decay"] = gr.Slider(value=0.0, minimum=-2.0, maximum=2.0, step=0.05, label="Repetition Penalty Length Decay", info="Modifies the reptition penalty based on how far back in time the token appeared in the sequence.") layout["inference_stt"]["inputs"]["length-penalty"] = gr.Slider(value=0.0, minimum=-2.0, maximum=2.0, step=0.05, label="Length Penalty", info="(AR only) Modifies the probability of a stop token based on the current length of the sequence.") with gr.Row():