vllm_omni.diffusion.models.ltx2.ltx2_transformer ¶

audio_sample: Tensor

sample: Tensor

approximate = approximate

proj = ColumnParallelLinear(
    dim_in,
    dim_out,
    bias=bias,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.proj" if prefix else "proj",
)

emb = PixArtAlphaCombinedTimestepSizeEmbeddings(
    embedding_dim,
    size_emb_dim=embedding_dim // 3,
    use_additional_conditions=use_additional_conditions,
)

linear = Linear(
    embedding_dim, num_mod_params * embedding_dim, bias=True
)

num_mod_params = num_mod_params

silu = SiLU()

attn = Attention(
    num_heads=query_num_heads,
    head_size=dim_head,
    num_kv_heads=kv_num_heads,
    softmax_scale=1.0 / dim_head**0.5,
    causal=False,
    prefix=prefix,
    disable_kv_quant=disable_kv_quant,
)

cross_attention_dim = (
    cross_attention_dim
    if cross_attention_dim is not None
    else query_dim
)

dropout = dropout

head_dim = dim_head

heads = query_num_heads

inner_dim = dim_head * heads

inner_kv_dim = dim_head * kv_heads

kv_num_heads = num_kv_heads

norm_k = TensorParallelRMSNorm(
    dim_head * kv_num_heads,
    eps=norm_eps,
    elementwise_affine=norm_elementwise_affine,
    tp_size=tp_size,
)

norm_q = TensorParallelRMSNorm(
    dim_head * query_num_heads,
    eps=norm_eps,
    elementwise_affine=norm_elementwise_affine,
    tp_size=tp_size,
)

out_dim = query_dim

query_dim = query_dim

query_num_heads = num_heads

rope_type = rope_type

to_gate_logits = Linear(query_dim, heads, bias=True)

to_k = None

to_out = ModuleList(
    [
        RowParallelLinear(
            inner_dim,
            out_dim,
            bias=out_bias,
            input_is_parallel=True,
            return_bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.to_out.0"
            if prefix
            else "to_out.0",
        ),
        Dropout(dropout) if dropout > 0 else Identity(),
    ]
)

to_q = None

to_qkv = None

to_v = None

total_num_heads = heads

total_num_kv_heads = kv_heads

use_bias = bias

audio_latents_per_second = (
    float(sampling_rate)
    / float(hop_length)
    / float(scale_factors[0])
)

base_height = base_height

base_num_frames = base_num_frames

base_width = base_width

causal_offset = causal_offset

dim = dim

double_precision = double_precision

hop_length = hop_length

modality = modality

num_attention_heads = num_attention_heads

patch_size = patch_size

patch_size_t = patch_size_t

rope_type = rope_type

sampling_rate = sampling_rate

scale_factors = scale_factors

theta = theta

net = ModuleList(layers)

audio_caption_projection = PixArtAlphaTextProjection(
    in_features=caption_channels,
    hidden_size=audio_inner_dim,
)

audio_norm_out = LayerNorm(
    audio_inner_dim, eps=1e-06, elementwise_affine=False
)

audio_proj_in = Linear(audio_in_channels, audio_inner_dim)

audio_proj_out = Linear(audio_inner_dim, audio_out_channels)

audio_prompt_adaln = LTX2AdaLayerNormSingle(
    audio_inner_dim,
    num_mod_params=2,
    use_additional_conditions=False,
)

audio_rope = LTX2AudioVideoRotaryPosEmbed(
    dim=audio_inner_dim,
    patch_size=audio_patch_size,
    patch_size_t=audio_patch_size_t,
    base_num_frames=audio_pos_embed_max_pos,
    sampling_rate=audio_sampling_rate,
    hop_length=audio_hop_length,
    scale_factors=[audio_scale_factor],
    theta=rope_theta,
    causal_offset=causal_offset,
    modality="audio",
    double_precision=rope_double_precision,
    rope_type=rope_type,
    num_attention_heads=audio_num_attention_heads,
)

audio_scale_shift_table = Parameter(
    randn(2, audio_inner_dim) / audio_inner_dim**0.5
)

audio_time_embed = LTX2AdaLayerNormSingle(
    audio_inner_dim,
    num_mod_params=audio_num_mod_params,
    use_additional_conditions=False,
)

av_cross_attn_audio_scale_shift = LTX2AdaLayerNormSingle(
    audio_inner_dim,
    num_mod_params=4,
    use_additional_conditions=False,
)

av_cross_attn_audio_v2a_gate = LTX2AdaLayerNormSingle(
    audio_inner_dim,
    num_mod_params=1,
    use_additional_conditions=False,
)

av_cross_attn_video_a2v_gate = LTX2AdaLayerNormSingle(
    inner_dim,
    num_mod_params=1,
    use_additional_conditions=False,
)

av_cross_attn_video_scale_shift = LTX2AdaLayerNormSingle(
    inner_dim,
    num_mod_params=4,
    use_additional_conditions=False,
)

caption_projection = PixArtAlphaTextProjection(
    in_features=caption_channels, hidden_size=inner_dim
)

config = SimpleNamespace(
    in_channels=in_channels,
    out_channels=out_channels,
    patch_size=patch_size,
    patch_size_t=patch_size_t,
    num_attention_heads=num_attention_heads,
    attention_head_dim=attention_head_dim,
    cross_attention_dim=cross_attention_dim,
    vae_scale_factors=vae_scale_factors,
    pos_embed_max_pos=pos_embed_max_pos,
    base_height=base_height,
    base_width=base_width,
    audio_in_channels=audio_in_channels,
    audio_out_channels=audio_out_channels,
    audio_patch_size=audio_patch_size,
    audio_patch_size_t=audio_patch_size_t,
    audio_num_attention_heads=audio_num_attention_heads,
    audio_attention_head_dim=audio_attention_head_dim,
    audio_cross_attention_dim=audio_cross_attention_dim,
    audio_scale_factor=audio_scale_factor,
    audio_pos_embed_max_pos=audio_pos_embed_max_pos,
    audio_sampling_rate=audio_sampling_rate,
    audio_hop_length=audio_hop_length,
    num_layers=num_layers,
    activation_fn=activation_fn,
    qk_norm=qk_norm,
    norm_elementwise_affine=norm_elementwise_affine,
    norm_eps=norm_eps,
    caption_channels=caption_channels,
    attention_bias=attention_bias,
    attention_out_bias=attention_out_bias,
    rope_theta=rope_theta,
    rope_double_precision=rope_double_precision,
    causal_offset=causal_offset,
    timestep_scale_multiplier=timestep_scale_multiplier,
    cross_attn_timestep_scale_multiplier=cross_attn_timestep_scale_multiplier,
    rope_type=rope_type,
)

cross_attn_audio_rope = LTX2AudioVideoRotaryPosEmbed(
    dim=audio_cross_attention_dim,
    patch_size=audio_patch_size,
    patch_size_t=audio_patch_size_t,
    base_num_frames=cross_attn_pos_embed_max_pos,
    sampling_rate=audio_sampling_rate,
    hop_length=audio_hop_length,
    theta=rope_theta,
    causal_offset=causal_offset,
    modality="audio",
    double_precision=rope_double_precision,
    rope_type=rope_type,
    num_attention_heads=audio_num_attention_heads,
)

cross_attn_rope = LTX2AudioVideoRotaryPosEmbed(
    dim=audio_cross_attention_dim,
    patch_size=patch_size,
    patch_size_t=patch_size_t,
    base_num_frames=cross_attn_pos_embed_max_pos,
    base_height=base_height,
    base_width=base_width,
    theta=rope_theta,
    causal_offset=causal_offset,
    modality="video",
    double_precision=rope_double_precision,
    rope_type=rope_type,
    num_attention_heads=num_attention_heads,
)

gradient_checkpointing = False

norm_out = LayerNorm(
    inner_dim, eps=1e-06, elementwise_affine=False
)

packed_modules_mapping = {
    "to_qkv": ["to_q", "to_k", "to_v"]
}

perturbed_attn = perturbed_attn

proj_in = Linear(in_channels, inner_dim)

proj_out = Linear(inner_dim, out_channels)

prompt_adaln = LTX2AdaLayerNormSingle(
    inner_dim,
    num_mod_params=2,
    use_additional_conditions=False,
)

prompt_modulation = cross_attn_mod or audio_cross_attn_mod

rope = LTX2AudioVideoRotaryPosEmbed(
    dim=inner_dim,
    patch_size=patch_size,
    patch_size_t=patch_size_t,
    base_num_frames=pos_embed_max_pos,
    base_height=base_height,
    base_width=base_width,
    scale_factors=vae_scale_factors,
    theta=rope_theta,
    causal_offset=causal_offset,
    modality="video",
    double_precision=rope_double_precision,
    rope_type=rope_type,
    num_attention_heads=num_attention_heads,
)

scale_shift_table = Parameter(
    randn(2, inner_dim) / inner_dim**0.5
)

time_embed = LTX2AdaLayerNormSingle(
    inner_dim,
    num_mod_params=video_num_mod_params,
    use_additional_conditions=False,
)

transformer_blocks = ModuleList(
    [
        (
            LTX2VideoTransformerBlock(
                dim=inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                cross_attention_dim=cross_attention_dim,
                audio_dim=audio_inner_dim,
                audio_num_attention_heads=audio_num_attention_heads,
                audio_attention_head_dim=audio_attention_head_dim,
                audio_cross_attention_dim=audio_cross_attention_dim,
                video_gated_attn=gated_attn,
                video_cross_attn_adaln=cross_attn_mod,
                audio_gated_attn=audio_gated_attn,
                audio_cross_attn_adaln=audio_cross_attn_mod,
                qk_norm=qk_norm,
                activation_fn=activation_fn,
                attention_bias=attention_bias,
                attention_out_bias=attention_out_bias,
                eps=norm_eps,
                elementwise_affine=norm_elementwise_affine,
                rope_type=rope_type,
                perturbed_attn=perturbed_attn,
                quant_config=quant_config,
                prefix=f"transformer_blocks.{layer_idx}",
            )
        )
        for layer_idx in (range(num_layers))
    ]
)

attn1 = LTX2Attention(
    query_dim=dim,
    heads=num_attention_heads,
    kv_heads=num_attention_heads,
    dim_head=attention_head_dim,
    bias=attention_bias,
    cross_attention_dim=None,
    out_bias=attention_out_bias,
    qk_norm=qk_norm,
    rope_type=rope_type,
    apply_gated_attention=video_gated_attn,
    quant_config=quant_config,
    prefix=f"{prefix}.attn1" if prefix else "attn1",
)

attn2 = LTX2Attention(
    query_dim=dim,
    cross_attention_dim=cross_attention_dim,
    heads=num_attention_heads,
    kv_heads=num_attention_heads,
    dim_head=attention_head_dim,
    bias=attention_bias,
    out_bias=attention_out_bias,
    qk_norm=qk_norm,
    rope_type=rope_type,
    apply_gated_attention=video_gated_attn,
    quant_config=quant_config,
    prefix=f"{prefix}.attn2" if prefix else "attn2",
    disable_kv_quant=True,
)

audio_a2v_cross_attn_scale_shift_table = Parameter(
    randn(5, audio_dim)
)

audio_attn1 = LTX2Attention(
    query_dim=audio_dim,
    heads=audio_num_attention_heads,
    kv_heads=audio_num_attention_heads,
    dim_head=audio_attention_head_dim,
    bias=attention_bias,
    cross_attention_dim=None,
    out_bias=attention_out_bias,
    qk_norm=qk_norm,
    rope_type=rope_type,
    apply_gated_attention=audio_gated_attn,
    quant_config=quant_config,
    prefix=f"{prefix}.audio_attn1"
    if prefix
    else "audio_attn1",
)

audio_attn2 = LTX2Attention(
    query_dim=audio_dim,
    cross_attention_dim=audio_cross_attention_dim,
    heads=audio_num_attention_heads,
    kv_heads=audio_num_attention_heads,
    dim_head=audio_attention_head_dim,
    bias=attention_bias,
    out_bias=attention_out_bias,
    qk_norm=qk_norm,
    rope_type=rope_type,
    apply_gated_attention=audio_gated_attn,
    quant_config=quant_config,
    prefix=f"{prefix}.audio_attn2"
    if prefix
    else "audio_attn2",
    disable_kv_quant=True,
)

audio_cross_attn_adaln = audio_cross_attn_adaln

audio_ff = LTX2FeedForward(
    audio_dim,
    activation_fn=activation_fn,
    quant_config=quant_config,
    prefix=f"{prefix}.audio_ff" if prefix else "audio_ff",
)

audio_norm1 = _make_rms_norm(
    audio_dim,
    eps=eps,
    elementwise_affine=elementwise_affine,
)

audio_norm2 = _make_rms_norm(
    audio_dim,
    eps=eps,
    elementwise_affine=elementwise_affine,
)

audio_norm3 = _make_rms_norm(
    audio_dim,
    eps=eps,
    elementwise_affine=elementwise_affine,
)

audio_prompt_scale_shift_table = Parameter(
    randn(2, audio_dim)
)

audio_scale_shift_table = Parameter(
    randn(audio_mod_param_num, audio_dim) / audio_dim**0.5
)

audio_to_video_attn = LTX2Attention(
    query_dim=dim,
    cross_attention_dim=audio_dim,
    heads=audio_num_attention_heads,
    kv_heads=audio_num_attention_heads,
    dim_head=audio_attention_head_dim,
    bias=attention_bias,
    out_bias=attention_out_bias,
    qk_norm=qk_norm,
    rope_type=rope_type,
    apply_gated_attention=video_gated_attn,
    quant_config=quant_config,
    prefix=f"{prefix}.audio_to_video_attn"
    if prefix
    else "audio_to_video_attn",
)

audio_to_video_norm = _make_rms_norm(
    dim, eps=eps, elementwise_affine=elementwise_affine
)

cross_attn_adaln = (
    video_cross_attn_adaln or audio_cross_attn_adaln
)

ff = LTX2FeedForward(
    dim,
    activation_fn=activation_fn,
    quant_config=quant_config,
    prefix=f"{prefix}.ff" if prefix else "ff",
)

norm1 = _make_rms_norm(
    dim, eps=eps, elementwise_affine=elementwise_affine
)

norm2 = _make_rms_norm(
    dim, eps=eps, elementwise_affine=elementwise_affine
)

norm3 = _make_rms_norm(
    dim, eps=eps, elementwise_affine=elementwise_affine
)

perturbed_attn = perturbed_attn

prompt_scale_shift_table = Parameter(randn(2, dim))

scale_shift_table = Parameter(
    randn(video_mod_param_num, dim) / dim**0.5
)

video_a2v_cross_attn_scale_shift_table = Parameter(
    randn(5, dim)
)

video_cross_attn_adaln = video_cross_attn_adaln

video_to_audio_attn = LTX2Attention(
    query_dim=audio_dim,
    cross_attention_dim=dim,
    heads=audio_num_attention_heads,
    kv_heads=audio_num_attention_heads,
    dim_head=audio_attention_head_dim,
    bias=attention_bias,
    out_bias=attention_out_bias,
    qk_norm=qk_norm,
    rope_type=rope_type,
    apply_gated_attention=audio_gated_attn,
    quant_config=quant_config,
    prefix=f"{prefix}.video_to_audio_attn"
    if prefix
    else "video_to_audio_attn",
)

video_to_audio_norm = _make_rms_norm(
    audio_dim,
    eps=eps,
    elementwise_affine=elementwise_affine,
)

eps = eps

global_hidden_size = hidden_size * max(tp_size, 1)

hidden_size = hidden_size

tp_size = tp_size

weight = Parameter(ones(hidden_size))