(
model_path,
spt_config_path,
spt_checkpoint_path,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
)
| 80 | |
| 81 | |
| 82 | def load_model( |
| 83 | model_path, |
| 84 | spt_config_path, |
| 85 | spt_checkpoint_path, |
| 86 | torch_dtype=torch.bfloat16, |
| 87 | attn_implementation="flash_attention_2", |
| 88 | ): |
| 89 | from transformers import AutoTokenizer |
| 90 | |
| 91 | from modeling_asteroid import AsteroidTTSInstruct |
| 92 | from XY_Tokenizer.xy_tokenizer.model import XY_Tokenizer |
| 93 | |
| 94 | tokenizer = AutoTokenizer.from_pretrained(model_path) |
| 95 | model = AsteroidTTSInstruct.from_pretrained( |
| 96 | model_path, torch_dtype=torch_dtype, attn_implementation=attn_implementation |
| 97 | ) |
| 98 | spt = XY_Tokenizer.load_from_checkpoint( |
| 99 | config_path=spt_config_path, ckpt_path=spt_checkpoint_path |
| 100 | ) |
| 101 | |
| 102 | model.eval() |
| 103 | spt.eval() |
| 104 | return tokenizer, model, spt |
| 105 | |
| 106 | |
| 107 | def process_jsonl_item(item): |
no test coverage detected