GPT_SoVITS/Accelerate/PyTorch/export.py

import enum
import os
import os.path as osp
import time
from pathlib import Path
from typing import MutableSequence, TypeAlias

import torch
import typer
from torch.export import Dim
from torch.nn import functional as F

from ..logger import logger
from . import nn
from .t2s_model_abc import AttentionABC, FeedForward, T2SDecoderABC, TransformerBlockABC, TransformerDecoderABC

Tensor = torch.Tensor

KVCache: TypeAlias = tuple[Tensor, Tensor]

app = typer.Typer(
    context_settings={"help_option_names": ["-h", "--help"]},
    add_completion=False,
)


class Stage(str, enum.Enum):
    embed = "embed"
    decode = "decode"


class KVCacheONNX:
    @staticmethod
    def empty(kv_cache):
        assert len(kv_cache) == 2
        k_cache, v_cache = kv_cache

        k_cache[:] = 0
        v_cache[:] = 0

    @staticmethod
    def update_cache(
        input_pos: Tensor, k_val: Tensor, v_val: Tensor, kv_cache: tuple[Tensor, Tensor], cache_idx: Tensor
    ):
        # input_pos: [B, ], k_val: [B, H, 1, D]
        k_out, v_out = kv_cache
        ip0 = input_pos - 1

        k_out[cache_idx, :, ip0, None] = k_val
        v_out[cache_idx, :, ip0, None] = v_val

        return k_out, v_out

    @staticmethod
    def prefill_kv(k_val: Tensor, v_val: Tensor, kv_cache: tuple[Tensor, Tensor]):
        # k_val: [B, S, H, D]
        k_cache, v_cache = kv_cache

        k_cache[..., : k_val.shape[1], :] = k_val.transpose(1, 2)
        v_cache[..., : v_val.shape[1], :] = v_val.transpose(1, 2)

    @staticmethod
    def init_cache(batch_size: int, max_seq_length: int, n_heads: int, head_dim: int, dtype: torch.dtype):
        cache_shape = (batch_size, n_heads, max_seq_length, head_dim)

        return (torch.zeros(cache_shape, dtype=dtype), torch.zeros(cache_shape, dtype=dtype))


class AttentionONNX(AttentionABC):
    def __init__(self, n_heads: int, head_dim: int, max_seq_length: int):
        super().__init__(n_heads, head_dim, max_seq_length)

        self.in_proj = nn.Linear(self.hidden_dim, self.hidden_dim * 3, bias=True)
        self.out_proj = nn.Linear(self.hidden_dim, self.hidden_dim, bias=True)

    def __call__(self, *args, **kwds):  # type: ignore
        pass

    def onnx_prefill(self, x: Tensor, kv_cache: KVCache, attn_mask: Tensor) -> Tensor:
        bsz, seqlen, _ = x.shape

        torch._check(attn_mask.size(-2) == x.size(-2))

        q, k, v = self.in_proj(x.unsqueeze(0)).chunk(3, dim=-1)

        q, k, v = map(lambda x: x.contiguous().view(bsz, seqlen, self.n_head, self.head_dim), (q, k, v))

        KVCacheONNX.prefill_kv(k, v, kv_cache)

        q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))

        attn = F.scaled_dot_product_attention(q, k, v, attn_mask)

        attn = attn.transpose(1, 2).contiguous().view(1, -1, self.hidden_dim)

        output = self.out_proj(attn)

        return output

    def onnx_decode(self, x: Tensor, input_pos: Tensor, kv_cache: KVCache, cache_idx: Tensor, attn_mask: Tensor):
        bsz, seqlen, _ = x.shape

        torch._check(attn_mask.size(-2) == 1)

        q, k, v = self.in_proj(x).chunk(3, dim=-1)

        q, k, v = map(lambda x: x.reshape(bsz, seqlen, self.n_head, self.head_dim), (q, k, v))

        q, k, v = map(lambda x: x.swapaxes(1, 2), (q, k, v))

        kv_cache = KVCacheONNX.update_cache(input_pos, k, v, kv_cache, cache_idx)

        max_idx = int(input_pos.max())

        q, k, v = map(lambda x: x[..., :max_idx, :], (q, *kv_cache))

        mask = attn_mask[..., :max_idx]

        attn = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)

        attn = attn.swapaxes(1, 2).reshape(bsz, seqlen, self.hidden_dim)

        attn = self.out_proj(attn)

        return attn


class TransformerBlockONNX(TransformerBlockABC):
    def __init__(self, n_head: int, ffn_dim: int, hidden_dim: int, max_seq_length: int) -> None:
        super().__init__(n_head, ffn_dim, hidden_dim, max_seq_length)

        self.attention: AttentionONNX = AttentionONNX(n_head, hidden_dim, max_seq_length)  # type: ignore
        self.feed_forward = FeedForward(hidden_dim, ffn_dim)
        self.attention_norm = nn.LayerNorm(self.hidden_dim)
        self.ffn_norm = nn.LayerNorm(self.hidden_dim)

    def onnx_prefill(self, x: Tensor, attn_mask: Tensor, kv_cache: KVCache):
        h = self.attention_norm(
            x
            + self.attention.onnx_prefill(
                x,
                kv_cache,
                attn_mask,
            )
        )
        out = self.ffn_norm(h + self.feed_forward(h))

        return out

    def onnx_decode(self, x: Tensor, input_pos: Tensor, kv_cache: KVCache, cache_idx: Tensor, attn_mask: Tensor):
        h = self.attention_norm(
            x
            + self.attention.onnx_decode(
                x,
                input_pos,
                kv_cache,
                cache_idx,
                attn_mask,
            )
        )
        out = self.ffn_norm(h + self.feed_forward(h))
        return out


class TransformerDecoderONNX(TransformerDecoderABC):
    def __init__(
        self,
        hidden_dim: int,
        n_layer: int,
        n_head: int,
        ffn_dim: int,
        vocab_size: int,
        max_seq_length: int,
        max_batch_size: int,
    ) -> None:
        super().__init__(hidden_dim, n_layer, n_head, ffn_dim, vocab_size, max_seq_length, max_batch_size)

        self.layers: MutableSequence[TransformerBlockONNX] = nn.ModuleList(  # type: ignore
            TransformerBlockONNX(n_head, ffn_dim, hidden_dim, max_seq_length) for _ in range(n_layer)
        )

    def onnx_prefill(self, x: Tensor, mask: Tensor, *kv_caches: KVCache):
        for layer, kv_cache in zip(self.layers, kv_caches):
            x = layer.onnx_prefill(
                x,
                mask,
                kv_cache,
            )
        return x

    def onnx_decode(
        self,
        input_pos: Tensor,
        x: Tensor,
        cache_idx: Tensor,
        attn_mask: Tensor,
        *kv_caches: KVCache,
    ):
        for layer, kv_cache in zip(self.layers, kv_caches):
            x = layer.onnx_decode(
                x,
                input_pos,
                kv_cache,
                cache_idx,
                attn_mask,
            )

        return x


class T2SDecoderONNX(T2SDecoderABC):
    def __init__(self, config: dict, max_seq_length: int = 2000, max_batch_size: int = 10) -> None:
        super().__init__(config, max_seq_length, max_batch_size)

        self.bert_proj = nn.Linear(1024, self.embedding_dim)
        self.ar_predict_layer = nn.Linear(self.hidden_dim, self.vocab_size, bias=False)

        self.h = TransformerDecoderONNX(
            self.hidden_dim, self.n_layer, self.n_head, self.ffn_dim, self.vocab_size, max_seq_length, max_batch_size
        )

    def pre_forward(self, session) -> tuple[list[Tensor], dict[str, Tensor]]:
        return super().pre_forward(session)

    def post_forward(self, idx: int, session) -> None:
        return super().post_forward(idx, session)

    def embed_onnx_(
        self,
        x: Tensor,
        x_len: Tensor,
        y: torch.Tensor,
        bert_features: Tensor,
    ):
        B = x.shape[0]
        D = self.embedding_dim
        T_TOTAL = 500
        xy_pos = torch.zeros((B, T_TOTAL, D)).to(bert_features[0].dtype)

        bert_features = bert_features.transpose(1, 2)

        y_len = y.shape[1]
        y_emb = self.ar_audio_embedding(y)
        y_pos = self.ar_audio_position.prefill(y_emb)

        for bs, x_, len_, bert_feature in zip(torch.arange(x.shape[0]), x, x_len, bert_features):
            x_emb = self.ar_text_embedding(x_[:len_])

            bert = self.bert_proj(bert_feature[:len_])

            print(bert.shape, bert_feature[:len_])

            return bert, bert_feature[:len_].unsqueeze(0)

            return bert[:20].unsqueeze(0), None
            x_emb = x_emb + bert
            x_pos = self.ar_text_position.prefill(x_emb.unsqueeze(0))

            xy_pos[None, bs, :len_] = bert
            # xy_pos[None, bs, len_ : len_ + y_len] = y_pos

        return xy_pos[:, -1], None

        return xy_pos[: x.shape[0]], x_len

    def embed_onnx(
        self,
        x: torch.Tensor,  # [B, Tx]
        x_len: torch.Tensor,  # [B]
        y: torch.Tensor,  # [1, Ty, D]
        bert_features: torch.Tensor,  # [B, 1024, Tx]
    ):
        # [B, 1024, Tx] -> [B, Tx, 1024]
        bert_features = bert_features.transpose(1, 2)

        Ty = y.shape[1]
        Tx = x.shape[1]
        B = x.shape[0]
        D = self.embedding_dim
        T_TOTAL = 500

        # mask: [B, Tx]，[j] Col < x_len[i]
        col = torch.arange(Tx, device=x.device).unsqueeze(0)  # [1, Tx]
        mask_x = col < x_len.view(-1, 1)  # [B, Tx]
        mask_x3 = mask_x.unsqueeze(-1)  # [B, Tx, 1]

        torch._check((Ty >= 0) and (Ty <= 250), "y_len out of range")
        torch._check((Tx >= 0) and (Tx <= 250), "x_len out of range")

        y_emb = self.ar_audio_embedding(y)  # [1, Ty, D]
        y_pos = self.ar_audio_position.prefill(y_emb)  # [1, Ty, D]

        x_emb_full = self.ar_text_embedding(x)  # [B, Tx, D]
        bert_full = self.bert_proj(bert_features[[0], : x_len[0]])  # [B, Tx, D]

        print(bert_full[0].shape, bert_features[0, : x_len[0]])

        return bert_full[0], bert_features[0, : x_len[0]]

        x_sum_full = x_emb_full + bert_full  # [B, Tx, D]
        x_pos_full = self.ar_text_position.prefill(x_sum_full)  # [B, Tx, D]

        xy_pos = torch.zeros((B, T_TOTAL, D), dtype=x_pos_full.dtype, device=x_pos_full.device)

        xy_pos[:, :Tx, :] = torch.where(
            mask_x3,
            bert_full[:, :Tx, :].to(xy_pos.dtype),
            xy_pos[:, :Tx, :],
        )

        return xy_pos[:, -1], None

        # Start From offset=x_len, Ty
        # [Ty] Index: offsets + [0..Ty-1]
        offsets = x_len.clamp(min=0, max=T_TOTAL - Ty)  # [B]
        idx_y = offsets.unsqueeze(1) + torch.arange(Ty, device=x_pos_full.device)  # [B, Ty]
        # scatter to dim=1
        # expand index to [B, Ty, D]
        idx_y3 = idx_y.unsqueeze(-1).expand(B, Ty, D)
        y_pos_b = y_pos.expand(B, Ty, D).to(xy_pos.dtype)  # [B, Ty, D]
        xy_pos = xy_pos.scatter(1, idx_y3, y_pos_b)

        return xy_pos, x_len


def torchscript_export(model: T2SDecoderONNX, stage="embed"):
    if stage == "embed":
        x = torch.randint(1, 600, (model.max_batch_size, 50))
        x_len = torch.randint(30, 50, (model.max_batch_size,))
        y = torch.randint(1, 600, (1, 200))
        bert_features = torch.rand((model.max_batch_size, 1024, 50))

        x_len[-1] = 50

        mask = torch.arange(x_len.max().item(), device=x.device).unsqueeze(0) < x_len.unsqueeze(1)

        x = x * mask
        bert_features = bert_features * mask.unsqueeze(1)

        try:
            a, c = model.embed_onnx_(x, x_len, y, bert_features)
            b, d = model.embed_onnx(x, x_len, y, bert_features)
            print("-" * 20)
            print(a - b, (a - b).sum(), (a - b).square().mean())
            print(c - d, (c - d).sum(), (c - d).square().mean())
            exit()
            assert torch.allclose(a, b, atol=1e-6, rtol=1e-8), (a - b).square().mean()

            setattr(model, "forward", model.embed_onnx)
            scripted_model = torch.jit.script(model, example_inputs=[(x, x_len, y, bert_features)])

            onnx_program = torch.onnx.export(
                scripted_model,
                (x, x_len, y, bert_features),
                input_names=["text", "text_len", "prompt", "bert_features"],
                output_names=["xy_pos", "input_pos"],
                dynamic_axes={
                    "text": {0: "Batch_Size", 1: "Sequence_Length_X"},
                    "prompt": {0: "Batch_Size", 1: "Sequence_Length_Y"},
                    "bert_features": {0: "Batch_Size", 1: "Sequence_Length_X"},
                },
                opset_version=21,
                training=False,
                do_constant_folding=True,
                external_data=False,
            )
            assert onnx_program
            onnx_program.save("onnx_export/AR_Embedding_TorchScript.onnx")

        except Exception:
            logger.bind(show_locals=False).exception("")


def dynamo_export(model: T2SDecoderONNX, stage="embed"):
    if stage == "embed":
        x = torch.randint(1, 600, (model.max_batch_size, 50))
        x_len = torch.randint(30, 50, (model.max_batch_size,))
        y = torch.randint(1, 600, (1, 200))
        bert_features = torch.rand((model.max_batch_size, 1024, 50))

        x_len[-1] = 50

        mask = torch.arange(x_len.max().item(), device=x.device).unsqueeze(0) < x_len.unsqueeze(1)

        x = x * mask
        bert_features = (bert_features.transpose(1, 2) * mask.unsqueeze(-1)).transpose(1, 2)

        dynamic_shapes = [
            {
                0: Dim("Batch_Size", min=1, max=4),
                1: Dim("Sequence_Length_X", min=1, max=50),
            },
            {
                0: Dim("Batch_Size", min=1, max=4),
            },
            {
                1: Dim("Sequence_Length_Y", min=1, max=250),
            },
            {
                0: Dim("Batch_Size", min=1, max=4),
                2: Dim("Sequence_Length_X", min=1, max=50),
            },
        ]
        try:
            a = model.embed_onnx_(x, x_len, y, bert_features)[0]
            b = model.embed_onnx(x, x_len, y, bert_features)[0]
            print(a - b, (a - b).square().mean())
            exit()
            assert torch.allclose(a, b, atol=1e-6, rtol=1e-8), (a - b).square().mean()

            setattr(model, "forward", model.embed_onnx)
            onnx_program = torch.onnx.export(
                model,
                (x, x_len, y, bert_features),
                input_names=["text", "text_len", "prompt", "bert_features"],
                output_names=["xy_pos", "input_pos"],
                dynamo=True,
                dynamic_shapes=dynamic_shapes,
                opset_version=21,
                training=False,
                do_constant_folding=True,
                external_data=False,
            )
            assert onnx_program
            onnx_program.save("onnx_export/AR_Embedding_Dynamo.onnx")
        except Exception:
            logger.bind(show_locals=False).exception("")


@app.command()
def export(
    ckpt_path: Path = typer.Option(
        ...,
        "--ckpt-path",
        file_okay=True,
        dir_okay=False,
        exists=True,
        readable=True,
        show_default=False,
        help="AR Checkpoint",
    ),
    dynamo: bool = typer.Option(False, is_flag=True, flag_value=True, help="Use Torch Dynamo"),
    stages: list[Stage] = typer.Option([Stage.embed], "--stages", help="Stage to export"),
):
    os.makedirs("onnx_export", exist_ok=True)
    dict_s1 = torch.load(ckpt_path, "cpu", mmap=True)
    condig = dict_s1["config"]
    model = T2SDecoderONNX(condig, 2000, 4)
    state_dict = dict_s1["weight"]
    model.load_state_dict(state_dict)

    for stage in stages:
        if dynamo:
            dynamo_export(model, stage)
        else:
            torchscript_export(model, stage)


def get_prog_name() -> str:
    script_rel = ".".join(["GPT_SoVITS", "Accelerate", "PyTorch", osp.basename(__file__)]).strip(".py")
    return f"python -s -m {script_rel}"


if __name__ == "__main__":
    t = time.perf_counter()
    app(prog_name=get_prog_name())
    logger.info(f"Exec Time: {time.perf_counter() - t:.2f} secs")