Function get_ltx2_transformer_config

scripts/convert_ltx2_to_diffusers.py:225–370

(version: str)

def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
    if version == "test":
        # Produces a transformer of the same size as used in test_models_transformer_ltx2.py
        config = {
            "model_id": "diffusers-internal-dev/dummy-ltx2",
            "diffusers_config": {
                "in_channels": 4,
                "out_channels": 4,
                "patch_size": 1,
                "patch_size_t": 1,
                "num_attention_heads": 2,
                "attention_head_dim": 8,
                "cross_attention_dim": 16,
                "vae_scale_factors": (8, 32, 32),
                "pos_embed_max_pos": 20,
                "base_height": 2048,
                "base_width": 2048,
                "audio_in_channels": 4,
                "audio_out_channels": 4,
                "audio_patch_size": 1,
                "audio_patch_size_t": 1,
                "audio_num_attention_heads": 2,
                "audio_attention_head_dim": 4,
                "audio_cross_attention_dim": 8,
                "audio_scale_factor": 4,
                "audio_pos_embed_max_pos": 20,
                "audio_sampling_rate": 16000,
                "audio_hop_length": 160,
                "num_layers": 2,
                "activation_fn": "gelu-approximate",
                "qk_norm": "rms_norm_across_heads",
                "norm_elementwise_affine": False,
                "norm_eps": 1e-6,
                "caption_channels": 16,
                "attention_bias": True,
                "attention_out_bias": True,
                "rope_theta": 10000.0,
                "rope_double_precision": False,
                "causal_offset": 1,
                "timestep_scale_multiplier": 1000,
                "cross_attn_timestep_scale_multiplier": 1,
            },
        }
        rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT
        special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP
    elif version == "2.0":
        config = {
            "model_id": "Lightricks/LTX-2",
            "diffusers_config": {
                "in_channels": 128,
                "out_channels": 128,
                "patch_size": 1,
                "patch_size_t": 1,
                "num_attention_heads": 32,
                "attention_head_dim": 128,
                "cross_attention_dim": 4096,
                "vae_scale_factors": (8, 32, 32),
                "pos_embed_max_pos": 20,