|
from typing import Generator |
|
|
|
import numpy as np |
|
import torch |
|
|
|
from modules import config, models |
|
from modules.utils.SeedContext import SeedContext |
|
|
|
|
|
@torch.inference_mode() |
|
def refine_text( |
|
text: str, |
|
prompt="[oral_2][laugh_0][break_6]", |
|
seed=-1, |
|
top_P=0.7, |
|
top_K=20, |
|
temperature=0.7, |
|
repetition_penalty=1.0, |
|
max_new_token=384, |
|
) -> str: |
|
chat_tts = models.load_chat_tts() |
|
|
|
with SeedContext(seed): |
|
refined_text = chat_tts.refiner_prompt( |
|
text, |
|
{ |
|
"prompt": prompt, |
|
"top_K": top_K, |
|
"top_P": top_P, |
|
"temperature": temperature, |
|
"repetition_penalty": repetition_penalty, |
|
"max_new_token": max_new_token, |
|
"disable_tqdm": config.runtime_env_vars.off_tqdm, |
|
}, |
|
) |
|
if isinstance(refined_text, Generator): |
|
raise NotImplementedError( |
|
"Refiner is not yet implemented for generator output" |
|
) |
|
if isinstance(refined_text, list): |
|
refined_text = "\n".join(refined_text) |
|
return refined_text |
|
|
