vllm_omni.diffusion.models.ltx2 ¶

component_profile = (
    LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
)

pipeline_kind = 'distilled_one_stage'

pipeline_recipe = LTX2_DISTILLED_ONE_STAGE_RECIPE

component_profile = LTX2_DISTILLED_COMPONENT_PROFILE

pipeline_kind = 'distilled_two_stage'

pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE

pipeline_kind = 'dmd2'

pipeline_recipe = LTX_POSITIVE_ONLY_RECIPE

unified_text_image_entry = False

component_profile = LTX2_COMPONENT_PROFILE

pipeline_kind = 'one_stage'

pipeline_recipe = LTX2_ONE_STAGE_RECIPE

supports_request_batch = True

unified_text_image_entry = True

pipeline_kind = 'dmd2'

pipeline_recipe = LTX_POSITIVE_ONLY_RECIPE

support_image_input = False

component_profile = LTX2_TWO_STAGE_COMPONENT_PROFILE

pipeline_kind = 'two_stage'

pipeline_recipe = LTX2_TWO_STAGE_RECIPE

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

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

audio_proj_in = nn.Linear(
    audio_in_channels, audio_inner_dim
)

audio_proj_out = nn.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 = nn.Parameter(
    torch.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,
    ff_bias=ff_bias,
    audio_ff_bias=audio_ff_bias,
    use_prompt_adaln_single=use_prompt_adaln_single,
    use_keyframes_abs_pos_embedding=use_keyframes_abs_pos_embedding,
)

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

keyframes_abs_pos_embedding = nn.Parameter(
    torch.zeros(1, inner_dim)
)

norm_out = nn.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 = nn.Linear(in_channels, inner_dim)

proj_out = nn.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 = nn.Parameter(
    torch.randn(2, inner_dim) / inner_dim**0.5
)

stacked_params_mapping = (
    (".attn1.to_qkv", ".attn1.to_q", "q"),
    (".attn1.to_qkv", ".attn1.to_k", "k"),
    (".attn1.to_qkv", ".attn1.to_v", "v"),
    (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"),
    (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"),
    (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"),
)

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

transformer_blocks = nn.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,
            ff_bias=ff_bias,
            audio_ff_bias=audio_ff_bias,
            quant_config=quant_config,
            prefix=f"transformer_blocks.{layer_idx}",
        )
        for layer_idx in range(num_layers)
    ]
)