Skip to content

vllm_omni.diffusion.models.soulx_singer.modules.llama

DiffLlama

Bases: LlamaModel

cond_mlp instance-attribute 

cond_mlp = nn.Sequential(
    nn.Linear(hidden_size, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, hidden_size),
)

diff_step_embedding instance-attribute 

diff_step_embedding = SinusoidalPosEmb(hidden_size)

diff_step_mlp instance-attribute 

diff_step_mlp = nn.Sequential(
    nn.Linear(hidden_size, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, hidden_size),
)

embed_tokens instance-attribute 

embed_tokens = None

layers instance-attribute 

layers = nn.ModuleList(
    [
        (LlamaNARDecoderLayer(layer_config, layer_idx=i))
        for i in (range(num_layers))
    ]
)

mel_mlp instance-attribute 

mel_mlp = nn.Sequential(
    nn.Linear(mel_dim, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, hidden_size),
)

mel_out_mlp instance-attribute 

mel_out_mlp = nn.Sequential(
    nn.Linear(hidden_size, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, mel_dim),
)

norm instance-attribute 

norm = LlamaAdaptiveRMSNorm(
    hidden_size, dim_cond=hidden_size
)

forward 

forward(
    x,
    diffusion_step,
    cond,
    x_mask,
    input_ids: LongTensor = None,
    attention_mask: Tensor | None = None,
    position_ids: LongTensor | None = None,
    past_key_values: list[FloatTensor] | None = None,
    inputs_embeds: FloatTensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = False,
) -> BaseModelOutputWithPast | Tensor | dict

LlamaAdaptiveRMSNorm

Bases: Module

to_weight instance-attribute 

to_weight = nn.Linear(dim_cond, hidden_size)

variance_epsilon instance-attribute 

variance_epsilon = eps

forward 

forward(hidden_states, cond_embedding)

LlamaNARDecoderLayer

Bases: LlamaDecoderLayer

input_layernorm instance-attribute 

input_layernorm = LlamaAdaptiveRMSNorm(
    config.hidden_size,
    eps=config.rms_norm_eps,
    dim_cond=config.hidden_size,
)

post_attention_layernorm instance-attribute 

post_attention_layernorm = LlamaAdaptiveRMSNorm(
    config.hidden_size,
    eps=config.rms_norm_eps,
    dim_cond=config.hidden_size,
)

forward 

forward(
    hidden_states: Tensor,
    cond_embedding: Tensor,
    attention_mask: Tensor | None = None,
    position_embeddings: tuple[Tensor, Tensor]
    | None = None,
    output_attentions: bool | None = False,
    cache_position: LongTensor | None = None,
    **kwargs,
) -> tuple[FloatTensor, ...]

SinusoidalPosEmb

Bases: Module

dim instance-attribute 

dim = dim

forward 

forward(x)