vllm_omni.diffusion.models.audiox.audiox_transformer ¶

attn = Attention(
    num_heads=local_nheads,
    head_size=head_dim,
    softmax_scale=head_dim**-0.5,
    causal=False,
)

dim = dim

head_dim = head_dim

k_norm = AudioXRMSNorm(head_dim)

nheads = nheads

q_norm = AudioXRMSNorm(head_dim)

to_kv = MergedColumnParallelLinear(
    input_size=dim,
    output_sizes=[dim, dim],
    bias=False,
    gather_output=False,
    prefix=f"{prefix}.to_kv",
)

to_q = ColumnParallelLinear(
    dim,
    dim,
    bias=False,
    gather_output=False,
    prefix=f"{prefix}.to_q",
)

w1 = _ColumnParallelChannelLastConv1d(
    dim,
    hidden_dim,
    bias=False,
    kernel_size=kernel_size,
    padding=padding,
)

w2 = _RowParallelChannelLastConv1d(
    hidden_dim,
    dim,
    bias=False,
    kernel_size=kernel_size,
    padding=padding,
)

w3 = _ColumnParallelChannelLastConv1d(
    dim,
    hidden_dim,
    bias=False,
    kernel_size=kernel_size,
    padding=padding,
)

adaLN_modulation = Sequential(
    SiLU(), Linear(dim, 6 * dim, bias=True)
)

attn = AudioXMMDiTSelfAttention(
    dim, nhead, prefix=f"{prefix}.attn"
)

cross_attn = AudioXCrossAttention(
    dim, nhead, prefix=f"{prefix}.cross_attn"
)

ffn = AudioXMMConvFeedForward(
    dim, int(dim * mlp_ratio), kernel_size=3, padding=1
)

linear1 = AudioXMMChannelLastConv1d(
    dim, dim, kernel_size=3, padding=1
)

norm1 = LayerNorm(dim, elementwise_affine=False)

norm2 = LayerNorm(dim, elementwise_affine=False)

attn = Attention(
    num_heads=num_heads,
    head_size=head_dim,
    softmax_scale=head_dim**-0.5,
    causal=False,
)

dim = dim

head_dim = head_dim

k_norm = AudioXRMSNorm(head_dim)

nheads = nheads

q_norm = AudioXRMSNorm(head_dim)

qkv = QKVParallelLinear(
    hidden_size=dim,
    head_size=head_dim,
    total_num_heads=nheads,
    bias=True,
    prefix=f"{prefix}.qkv",
)

rope = RotaryEmbedding(is_neox_style=False)

eps = eps

depth = depth

device

dim = dim

layers = ModuleList(
    [
        (
            AudioXMMDiTBlock(
                hidden_dim,
                num_heads,
                mlp_ratio=mlp_ratio,
                prefix=f"layers.{i}",
            )
        )
        for i in (range(depth))
    ]
)

proj_mm_seq_len = (
    Linear(384, _latent_seq_len)
    if _latent_seq_len != 384
    else Identity()
)

proj_mm_tokens = (
    Linear(768, hidden_dim) if dim != 768 else Identity()
)

project_in = (
    Linear(dim_in, dim, bias=False)
    if dim_in is not None
    else Identity()
)

project_out = (
    Linear(dim, dim_out, bias=False)
    if dim_out is not None
    else Identity()
)

cond_token_dim = cond_token_dim

postprocess_conv = Conv1d(
    io_channels, io_channels, 1, bias=False
)

preprocess_conv = Conv1d(
    io_channels, io_channels, 1, bias=False
)

timestep_features = GaussianFourierProjection(
    in_features=1,
    embedding_size=timestep_features_dim // 2,
    scale=1.0,
    trainable=False,
)

to_cond_embed = Sequential(
    Linear(cond_token_dim, cond_embed_dim, bias=False),
    SiLU(),
    Linear(cond_embed_dim, cond_embed_dim, bias=False),
)

to_global_embed = Sequential(
    Linear(global_cond_dim, global_embed_dim, bias=False),
    SiLU(),
    Linear(global_embed_dim, global_embed_dim, bias=False),
)

to_timestep_embed = Sequential(
    Linear(timestep_features_dim, embed_dim, bias=True),
    SiLU(),
    Linear(embed_dim, embed_dim, bias=True),
)

transformer = ContinuousMMDiTTransformer(
    dim=embed_dim,
    depth=depth,
    dim_heads=embed_dim // num_heads,
    dim_in=io_channels,
    dim_out=io_channels,
)