mirror of
https://github.com/2noise/ChatTTS.git
synced 2026-08-29 02:10:59 +08:00
feat: add examples of onnx export (#622)
This commit is contained in:
ZillaRU
@@ -0,0 +1,14 @@
|
||||
# Export onnx or JIT models for deployment
|
||||
|
||||
## Run `pip install onnx -U`.
|
||||
|
||||
## Export GPT
|
||||
|
||||
3. Run `python examples/onnx/exporter.py --gpt`
|
||||
|
||||
|
||||
## Export other models
|
||||
Run `python examples/onnx/exporter.py --decoder --vocos`
|
||||
|
||||
## Reference
|
||||
[Run LLMs on Sophon TPU](https://github.com/sophgo/LLM-TPU)
|
||||
@@ -0,0 +1,417 @@
|
||||
from dataclasses import asdict
|
||||
import argparse
|
||||
import os
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from ChatTTS.model.dvae import DVAE
|
||||
from ChatTTS.config import Config
|
||||
from vocos import Vocos
|
||||
from vocos.pretrained import instantiate_class
|
||||
import torch.jit as jit
|
||||
from examples.onnx.gpt import GPT
|
||||
|
||||
# disable cuda
|
||||
torch.cuda.is_available = lambda : False
|
||||
|
||||
# add args to control which modules to export
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--gpt", action="store_true", help="trace gpt")
|
||||
parser.add_argument("--decoder", action="store_true", help="trace decoder")
|
||||
parser.add_argument("--vocos", action="store_true", help="trace vocos")
|
||||
parser.add_argument("--pth_dir", default="./assets", type=str, help="path to the pth model directory")
|
||||
parser.add_argument("--out_dir", default="./tmp", type=str, help="path to output directory")
|
||||
|
||||
args = parser.parse_args()
|
||||
chattts_config = Config()
|
||||
|
||||
def export_gpt():
|
||||
gpt_model = GPT(
|
||||
gpt_config=asdict(chattts_config.gpt),
|
||||
use_flash_attn=False
|
||||
).eval()
|
||||
gpt_model.from_pretrained(asdict(chattts_config.path)['gpt_ckpt_path'])
|
||||
gpt_model = gpt_model.eval()
|
||||
for param in gpt_model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
config = gpt_model.gpt.config
|
||||
layers = gpt_model.gpt.layers
|
||||
model_norm = gpt_model.gpt.norm
|
||||
|
||||
NUM_OF_LAYERS = config.num_hidden_layers
|
||||
HIDDEN_SIZE = config.hidden_size
|
||||
NUM_ATTENTION_HEADS = config.num_attention_heads
|
||||
NUM_KEY_VALUE_HEADS = config.num_key_value_heads
|
||||
HEAD_DIM = HIDDEN_SIZE // NUM_ATTENTION_HEADS # 64
|
||||
TEXT_VOCAB_SIZE = gpt_model.emb_text.weight.shape[0]
|
||||
AUDIO_VOCAB_SIZE = gpt_model.emb_code[0].weight.shape[0]
|
||||
SEQ_LENGTH = 512
|
||||
|
||||
folder = os.path.join(args.out_dir, "gpt")
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
|
||||
for param in gpt_model.emb_text.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
for param in gpt_model.emb_code.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
for param in gpt_model.head_code.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
for param in gpt_model.head_text.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
class EmbeddingText(torch.nn.Module):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def forward(self, input_ids):
|
||||
return gpt_model.emb_text(input_ids)
|
||||
|
||||
def convert_embedding_text():
|
||||
model = EmbeddingText()
|
||||
input_ids = torch.tensor([range(SEQ_LENGTH)])
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(input_ids),
|
||||
f"{folder}/embedding_text.onnx",
|
||||
verbose=False,
|
||||
input_names=["input_ids"],
|
||||
output_names=["input_embed"],
|
||||
do_constant_folding=True,
|
||||
opset_version=15,
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingCode(torch.nn.Module):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def forward(self, input_ids):
|
||||
input_ids = input_ids.unsqueeze(2).expand(
|
||||
-1, -1, gpt_model.num_vq
|
||||
) # for forward_first_code
|
||||
code_emb = [
|
||||
gpt_model.emb_code[i](input_ids[:, :, i]) for i in range(gpt_model.num_vq)
|
||||
]
|
||||
return torch.stack(code_emb, 2).sum(2)
|
||||
|
||||
|
||||
def convert_embedding_code():
|
||||
model = EmbeddingCode()
|
||||
input_ids = torch.tensor([range(SEQ_LENGTH)])
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(input_ids),
|
||||
f"{folder}/embedding_code.onnx",
|
||||
verbose=False,
|
||||
input_names=["input_ids"],
|
||||
output_names=["input_embed"],
|
||||
do_constant_folding=True,
|
||||
opset_version=15,
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingCodeCache(torch.nn.Module): # for forward_next_code
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def forward(self, input_ids):
|
||||
code_emb = [
|
||||
gpt_model.emb_code[i](input_ids[:, :, i]) for i in range(gpt_model.num_vq)
|
||||
]
|
||||
return torch.stack(code_emb, 2).sum(2)
|
||||
|
||||
|
||||
def convert_embedding_code_cache():
|
||||
model = EmbeddingCodeCache()
|
||||
input_ids = torch.tensor(
|
||||
[[[416, 290, 166, 212]]]
|
||||
) # torch.tensor([[range(gpt_model.num_vq)]])
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(input_ids),
|
||||
f"{folder}/embedding_code_cache.onnx",
|
||||
verbose=False,
|
||||
input_names=["input_ids"],
|
||||
output_names=["input_embed"],
|
||||
do_constant_folding=True,
|
||||
opset_version=15,
|
||||
)
|
||||
|
||||
|
||||
class Block(torch.nn.Module):
|
||||
def __init__(self, layer_id):
|
||||
super().__init__()
|
||||
self.layer_id = layer_id
|
||||
self.layer = layers[layer_id] # LlamaDecoderLayer
|
||||
self.norm = model_norm
|
||||
|
||||
def forward(self, hidden_states, position_ids, attention_mask):
|
||||
hidden_states, past_kv = self.layer(
|
||||
hidden_states=hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
use_cache=True,
|
||||
)
|
||||
present_k, present_v = past_kv
|
||||
if self.layer_id == NUM_OF_LAYERS - 1:
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return hidden_states, present_k, present_v
|
||||
|
||||
|
||||
def convert_block(layer_id):
|
||||
model = Block(layer_id)
|
||||
hidden_states = torch.randn((1, SEQ_LENGTH, HIDDEN_SIZE))
|
||||
position_ids = torch.tensor([range(SEQ_LENGTH)], dtype=torch.long)
|
||||
attention_mask = -1000 * torch.ones(
|
||||
(1, 1, SEQ_LENGTH, SEQ_LENGTH), dtype=torch.float32
|
||||
).triu(diagonal=1)
|
||||
model(hidden_states, position_ids, attention_mask)
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(hidden_states, position_ids, attention_mask),
|
||||
f"{folder}/block_{layer_id}.onnx",
|
||||
verbose=False,
|
||||
input_names=["input_states", "position_ids", "attention_mask"],
|
||||
output_names=["hidden_states", "past_k", "past_v"],
|
||||
do_constant_folding=True,
|
||||
opset_version=15,
|
||||
)
|
||||
|
||||
|
||||
class BlockCache(torch.nn.Module):
|
||||
|
||||
def __init__(self, layer_id):
|
||||
super().__init__()
|
||||
self.layer_id = layer_id
|
||||
self.layer = layers[layer_id]
|
||||
self.norm = model_norm
|
||||
|
||||
def forward(self, hidden_states, position_ids, attention_mask, past_k, past_v):
|
||||
hidden_states, past_kv = self.layer(
|
||||
hidden_states,
|
||||
attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=(past_k, past_v),
|
||||
use_cache=True,
|
||||
)
|
||||
present_k, present_v = past_kv
|
||||
if self.layer_id == NUM_OF_LAYERS - 1:
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return hidden_states, present_k, present_v
|
||||
|
||||
|
||||
def convert_block_cache(layer_id):
|
||||
model = BlockCache(layer_id)
|
||||
hidden_states = torch.randn((1, 1, HIDDEN_SIZE))
|
||||
position_ids = torch.tensor([range(1)], dtype=torch.long)
|
||||
attention_mask = -1000 * torch.ones(
|
||||
(1, 1, 1, SEQ_LENGTH + 1), dtype=torch.float32
|
||||
).triu(diagonal=1)
|
||||
past_k = torch.randn((1, SEQ_LENGTH, NUM_ATTENTION_HEADS, HEAD_DIM))
|
||||
past_v = torch.randn((1, SEQ_LENGTH, NUM_ATTENTION_HEADS, HEAD_DIM))
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(hidden_states, position_ids, attention_mask, past_k, past_v),
|
||||
f"{folder}/block_cache_{layer_id}.onnx",
|
||||
verbose=False,
|
||||
input_names=[
|
||||
"input_states",
|
||||
"position_ids",
|
||||
"attention_mask",
|
||||
"history_k",
|
||||
"history_v",
|
||||
],
|
||||
output_names=["hidden_states", "past_k", "past_v"],
|
||||
do_constant_folding=True,
|
||||
opset_version=15,
|
||||
)
|
||||
|
||||
|
||||
class GreedyHead(torch.nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, m_logits):
|
||||
_, token = torch.topk(m_logits.float(), 1)
|
||||
return token
|
||||
|
||||
|
||||
def convert_greedy_head_text():
|
||||
model = GreedyHead()
|
||||
m_logits = torch.randn(1, TEXT_VOCAB_SIZE)
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(m_logits),
|
||||
f"{folder}/greedy_head_text.onnx",
|
||||
verbose=False,
|
||||
input_names=["m_logits"],
|
||||
output_names=["token"],
|
||||
do_constant_folding=True,
|
||||
opset_version=15,
|
||||
)
|
||||
|
||||
|
||||
def convert_greedy_head_code():
|
||||
model = GreedyHead()
|
||||
m_logits = torch.randn(1, AUDIO_VOCAB_SIZE, gpt_model.num_vq)
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(m_logits),
|
||||
f"{folder}/greedy_head_code.onnx",
|
||||
verbose=False,
|
||||
input_names=["m_logits"],
|
||||
output_names=["token"],
|
||||
do_constant_folding=True,
|
||||
opset_version=15,
|
||||
)
|
||||
|
||||
class LmHead_infer_text(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, hidden_states):
|
||||
m_logits = gpt_model.head_text(hidden_states)
|
||||
return m_logits
|
||||
|
||||
|
||||
class LmHead_infer_code(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, hidden_states):
|
||||
m_logits = torch.stack(
|
||||
[gpt_model.head_code[i](hidden_states) for i in range(gpt_model.num_vq)], 2
|
||||
)
|
||||
return m_logits
|
||||
|
||||
|
||||
def convert_lm_head_text():
|
||||
model = LmHead_infer_text()
|
||||
input = torch.randn(1, HIDDEN_SIZE)
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(input),
|
||||
f"{folder}/lm_head_text.onnx",
|
||||
verbose=False,
|
||||
input_names=["hidden_states"],
|
||||
output_names=["m_logits"],
|
||||
do_constant_folding=True,
|
||||
opset_version=15,
|
||||
)
|
||||
|
||||
|
||||
def convert_lm_head_code():
|
||||
model = LmHead_infer_code()
|
||||
input = torch.randn(1, HIDDEN_SIZE)
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(input),
|
||||
f"{folder}/lm_head_code.onnx",
|
||||
verbose=False,
|
||||
input_names=["hidden_states"],
|
||||
output_names=["m_logits"],
|
||||
do_constant_folding=True,
|
||||
opset_version=15,
|
||||
)
|
||||
|
||||
# export models
|
||||
print(f"Convert block & block_cache")
|
||||
for i in tqdm(range(NUM_OF_LAYERS)):
|
||||
convert_block(i)
|
||||
convert_block_cache(i)
|
||||
|
||||
print(f"Convert embedding")
|
||||
convert_embedding_text()
|
||||
convert_embedding_code()
|
||||
convert_embedding_code_cache()
|
||||
|
||||
print(f"Convert lm_head")
|
||||
convert_lm_head_code()
|
||||
convert_lm_head_text()
|
||||
|
||||
print(f"Convert greedy_head")
|
||||
convert_greedy_head_text()
|
||||
convert_greedy_head_code()
|
||||
|
||||
def export_decoder():
|
||||
decoder = (
|
||||
DVAE(
|
||||
decoder_config=asdict(chattts_config.decoder),
|
||||
dim=chattts_config.decoder.idim,
|
||||
).eval()
|
||||
)
|
||||
decoder.load_state_dict(
|
||||
torch.load(asdict(chattts_config.path)['decoder_ckpt_path'], weights_only=True, mmap=True)
|
||||
)
|
||||
|
||||
for param in decoder.parameters():
|
||||
param.requires_grad = False
|
||||
rand_input = torch.rand([1, 768, 1024], requires_grad=False)
|
||||
|
||||
def mydec(_inp):
|
||||
return decoder(_inp, mode='decode')
|
||||
|
||||
jitmodel = jit.trace(mydec, [rand_input])
|
||||
jit.save(jitmodel, f"{args.out_dir}/decoder_jit.pt")
|
||||
|
||||
|
||||
def export_vocos():
|
||||
feature_extractor = instantiate_class(
|
||||
args=(), init=asdict(chattts_config.vocos.feature_extractor)
|
||||
)
|
||||
backbone = instantiate_class(args=(), init=asdict(chattts_config.vocos.backbone))
|
||||
head = instantiate_class(args=(), init=asdict(chattts_config.vocos.head))
|
||||
vocos = (
|
||||
Vocos(feature_extractor=feature_extractor, backbone=backbone, head=head).eval()
|
||||
)
|
||||
vocos.load_state_dict(torch.load(asdict(chattts_config.path)['vocos_ckpt_path'], weights_only=True, mmap=True))
|
||||
|
||||
for param in vocos.parameters():
|
||||
param.requires_grad = False
|
||||
rand_input = torch.rand([1, 100, 2048], requires_grad=False)
|
||||
|
||||
def myvocos(_inp):
|
||||
# return chat.vocos.decode(_inp) # TPU cannot support the istft OP, thus it has to be moved to postprocessing
|
||||
# reference: https://github.com/gemelo-ai/vocos.git
|
||||
x = vocos.backbone(_inp)
|
||||
x = vocos.head.out(x).transpose(1, 2)
|
||||
mag, p = x.chunk(2, dim=1)
|
||||
mag = torch.exp(mag)
|
||||
mag = torch.clip(
|
||||
mag, max=1e2
|
||||
) # safeguard to prevent excessively large magnitudes
|
||||
# wrapping happens here. These two lines produce real and imaginary value
|
||||
x = torch.cos(p)
|
||||
y = torch.sin(p)
|
||||
return mag, x, y
|
||||
|
||||
jitmodel = jit.trace(myvocos, [rand_input])
|
||||
torch.onnx.export(
|
||||
jitmodel,
|
||||
[rand_input],
|
||||
f"{args.out_dir}/vocos_1-100-2048.onnx",
|
||||
opset_version=12,
|
||||
do_constant_folding=True,
|
||||
)
|
||||
|
||||
|
||||
if args.gpt:
|
||||
export_gpt()
|
||||
|
||||
if args.decoder:
|
||||
export_decoder()
|
||||
|
||||
if args.vocos:
|
||||
export_vocos()
|
||||
|
||||
print("Done. Please check the files in", args.out_dir)
|
||||
@@ -0,0 +1,85 @@
|
||||
import logging
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn.utils.parametrizations import weight_norm
|
||||
from .modeling_llama import LlamaModel, LlamaConfig
|
||||
|
||||
class GPT(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
gpt_config: dict,
|
||||
num_audio_tokens: int = 626,
|
||||
num_text_tokens: int = 21178,
|
||||
num_vq=4,
|
||||
use_flash_attn=False,
|
||||
device=torch.device("cpu"),
|
||||
logger=logging.getLogger(__name__),
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.logger = logger
|
||||
|
||||
self.device = device
|
||||
self.device_gpt = device if "mps" not in str(device) else torch.device("cpu")
|
||||
|
||||
self.num_vq = num_vq
|
||||
self.num_audio_tokens = num_audio_tokens
|
||||
|
||||
self.use_flash_attn = use_flash_attn
|
||||
|
||||
self.gpt, self.llama_config = self._build_llama(gpt_config, self.device_gpt)
|
||||
self.is_te_llama = False
|
||||
self.model_dim = int(self.gpt.config.hidden_size)
|
||||
self.emb_code = nn.ModuleList(
|
||||
[
|
||||
nn.Embedding(
|
||||
num_audio_tokens,
|
||||
self.model_dim,
|
||||
device=self.device_gpt,
|
||||
)
|
||||
for _ in range(num_vq)
|
||||
],
|
||||
)
|
||||
self.emb_text = nn.Embedding(
|
||||
num_text_tokens, self.model_dim, device=self.device_gpt
|
||||
)
|
||||
|
||||
self.head_text = weight_norm(
|
||||
nn.Linear(
|
||||
self.model_dim,
|
||||
num_text_tokens,
|
||||
bias=False,
|
||||
device=device,
|
||||
),
|
||||
name="weight",
|
||||
)
|
||||
self.head_code = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
nn.Linear(
|
||||
self.model_dim,
|
||||
num_audio_tokens,
|
||||
bias=False,
|
||||
device=device,
|
||||
),
|
||||
name="weight",
|
||||
)
|
||||
for _ in range(self.num_vq)
|
||||
],
|
||||
)
|
||||
|
||||
def from_pretrained(self, file_path: str):
|
||||
self.load_state_dict(torch.load(file_path, weights_only=True, mmap=True), strict=False)
|
||||
|
||||
def _build_llama(
|
||||
self,
|
||||
config: dict,
|
||||
device: torch.device,
|
||||
) -> Tuple[LlamaModel, LlamaConfig]:
|
||||
|
||||
llama_config = LlamaConfig(**config)
|
||||
model = LlamaModel(llama_config)
|
||||
del model.embed_tokens
|
||||
return model.to(device), llama_config
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user