vllm_omni.model_executor.models.audio8_tts.codec ¶

conv = nn.Conv1d(
    in_channels,
    out_channels,
    kernel_size,
    stride=stride,
    dilation=dilation,
    groups=groups,
)

kernel_size = (kernel_size - 1) * dilation + 1

padding = self.kernel_size - self.stride

stride = stride

conv = nn.ConvTranspose1d(
    in_channels,
    out_channels,
    kernel_size,
    stride=stride,
    dilation=dilation,
)

kernel_size = kernel_size

stride = stride

decoder = ArkttsDecoder()

encoder = ArkttsEncoder()

frame_length = 2048

quantizer = ArkttsDownsampleQuantizer(
    post_n_layer=post_n_layer,
    post_n_head=post_n_head,
    post_n_local_heads=post_n_local_heads,
    post_intermediate_size=post_intermediate_size,
)

sample_rate = 44100

head_dim = config.head_dim

n_head = config.n_head

n_local_heads = config.n_local_heads

wo = nn.Linear(
    config.head_dim * config.n_head, config.dim, bias=False
)

wqkv = nn.Linear(config.dim, total, bias=False)

w1 = nn.Linear(
    config.dim, config.intermediate_size, bias=False
)

w2 = nn.Linear(
    config.intermediate_size, config.dim, bias=False
)

w3 = nn.Linear(
    config.dim, config.intermediate_size, bias=False
)

gamma = nn.Parameter(init_values * torch.ones(dim))

eps = float(eps)

weight = nn.Parameter(torch.ones(dim))

attention = ArkttsCodecAttention(config)

attention_layer_scale = ArkttsCodecLayerScale(config.dim)

attention_norm = ArkttsCodecRMSNorm(
    config.dim, config.norm_eps
)

feed_forward = ArkttsCodecFeedForward(config)

ffn_layer_scale = ArkttsCodecLayerScale(config.dim)

ffn_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)

channels_first = bool(channels_first)

dim = int(dim)

head_dim = int(head_dim)

intermediate_size = int(intermediate_size)

n_head = int(n_head)

n_layer = int(n_layer)

n_local_heads = (
    int(n_head)
    if n_local_heads == -1
    else int(n_local_heads)
)

norm_eps = float(norm_eps)

rope_base = float(rope_base)

channels_first = config.channels_first

head_dim = config.head_dim

input_proj = (
    nn.Linear(input_dim, config.dim)
    if input_dim != config.dim
    else nn.Identity()
)

layers = nn.ModuleList(
    [
        ArkttsCodecTransformerBlock(config)
        for _ in range(config.n_layer)
    ]
)

look_ahead_conv = nn.Identity()

norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)

output_proj = (
    nn.Linear(config.dim, input_dim)
    if input_dim != config.dim
    else nn.Identity()
)

rope_base = config.rope_base

window_size = window_size

act = nn.GELU()

dwconv = ArkttsCausalConv1d(
    dim, dim, kernel_size=7, groups=dim
)

gamma = nn.Parameter(1e-06 * torch.ones(dim))

norm = nn.LayerNorm(dim, eps=1e-06)

pwconv1 = nn.Linear(dim, 4 * dim)

pwconv2 = nn.Linear(4 * dim, dim)

model = nn.Sequential(*modules)

block = nn.Sequential(
    ArkttsSnake1d(input_dim),
    _causal_wn_transpose(
        input_dim,
        output_dim,
        kernel_size=2 * stride,
        stride=stride,
    ),
    ArkttsResidualUnit(output_dim, 1),
    ArkttsResidualUnit(output_dim, 3),
    ArkttsResidualUnit(output_dim, 9),
)

downsample = nn.Sequential(
    nn.Sequential(
        ArkttsCausalConv1d(
            1024, 1024, kernel_size=2, stride=2
        ),
        ArkttsConvNeXtBlock(1024),
    ),
    nn.Sequential(
        ArkttsCausalConv1d(
            1024, 1024, kernel_size=2, stride=2
        ),
        ArkttsConvNeXtBlock(1024),
    ),
)

post_module = ArkttsCodecWindowTransformer(
    ArkttsCodecTransformerConfig(
        n_layer=post_n_layer,
        n_head=post_n_head,
        n_local_heads=post_n_local_heads,
        dim=1024,
        intermediate_size=post_intermediate_size,
    ),
    1024,
    window_size=128,
)

pre_module = ArkttsCodecWindowTransformer(
    ArkttsCodecTransformerConfig(
        n_layer=8,
        n_head=16,
        dim=1024,
        intermediate_size=3072,
    ),
    1024,
    window_size=128,
)

quantizer = ArkttsResidualQuantizer(1024, 9, 1024, 8)

semantic_predictor_module = nn.Identity()

semantic_quantizer = ArkttsResidualQuantizer(
    1024, 1, 4096, 8
)

upsample = nn.Sequential(
    nn.Sequential(
        ArkttsCausalConvTranspose1d(
            1024, 1024, kernel_size=2, stride=2
        ),
        ArkttsConvNeXtBlock(1024),
    ),
    nn.Sequential(
        ArkttsCausalConvTranspose1d(
            1024, 1024, kernel_size=2, stride=2
        ),
        ArkttsConvNeXtBlock(1024),
    ),
)

block = nn.Sequential(*modules)

block = nn.Sequential(*modules)

codebook_size = int(codebook_size)

n_codebooks = int(n_codebooks)

quantizers = nn.ModuleList(
    [
        ArkttsVectorQuantizer(
            input_dim, codebook_size, codebook_dim
        )
        for _ in range(n_codebooks)
    ]
)

block = nn.Sequential(
    ArkttsSnake1d(dim),
    _causal_wn_conv(
        dim, dim, kernel_size=7, dilation=dilation
    ),
    ArkttsSnake1d(dim),
    _causal_wn_conv(dim, dim, kernel_size=1),
)

alpha = nn.Parameter(torch.ones(1, channels, 1))

codebook = nn.Embedding(codebook_size, codebook_dim)

codebook_dim = int(codebook_dim)

codebook_size = int(codebook_size)

in_proj = legacy_weight_norm(
    nn.Conv1d(input_dim, codebook_dim, kernel_size=1)
)

out_proj = legacy_weight_norm(
    nn.Conv1d(codebook_dim, input_dim, kernel_size=1)
)