| 1 | import torch |
| 2 | |
| 3 | from ltx_core.loader.sd_ops import SDOps |
| 4 | from ltx_core.model.model_protocol import ModelConfigurator |
| 5 | from ltx_core.model.transformer.attention import AttentionFunction |
| 6 | from ltx_core.model.transformer.model import LTXModel, LTXModelType |
| 7 | from ltx_core.model.transformer.rope import LTXRopeType |
| 8 | from ltx_core.model.transformer.text_projection import create_caption_projection |
| 9 | from ltx_core.utils import check_config_value |
| 10 | |
| 11 | |
| 12 | class LTXModelConfigurator(ModelConfigurator[LTXModel]): |
| 13 | """ |
| 14 | Configurator for LTX model. |
| 15 | Used to create an LTX model from a configuration dictionary. |
| 16 | """ |
| 17 | |
| 18 | @classmethod |
| 19 | def from_config(cls: type[LTXModel], config: dict) -> LTXModel: |
| 20 | # Build caption projections for 19B models (projection handled in transformer). |
| 21 | caption_projection, audio_caption_projection = _build_caption_projections(config, is_av=True) |
| 22 | |
| 23 | config = config.get("transformer", {}) |
| 24 | |
| 25 | check_config_value(config, "dropout", 0.0) |
| 26 | check_config_value(config, "attention_bias", True) |
| 27 | check_config_value(config, "num_vector_embeds", None) |
| 28 | check_config_value(config, "activation_fn", "gelu-approximate") |
| 29 | check_config_value(config, "num_embeds_ada_norm", 1000) |
| 30 | check_config_value(config, "use_linear_projection", False) |
| 31 | check_config_value(config, "only_cross_attention", False) |
| 32 | check_config_value(config, "cross_attention_norm", True) |
| 33 | check_config_value(config, "double_self_attention", False) |
| 34 | check_config_value(config, "upcast_attention", False) |
| 35 | check_config_value(config, "standardization_norm", "rms_norm") |
| 36 | check_config_value(config, "norm_elementwise_affine", False) |
| 37 | check_config_value(config, "qk_norm", "rms_norm") |
| 38 | check_config_value(config, "positional_embedding_type", "rope") |
| 39 | check_config_value(config, "use_audio_video_cross_attention", True) |
| 40 | check_config_value(config, "share_ff", False) |
| 41 | check_config_value(config, "av_cross_ada_norm", True) |
| 42 | check_config_value(config, "use_middle_indices_grid", True) |
| 43 | |
| 44 | return LTXModel( |
| 45 | model_type=LTXModelType.AudioVideo, |
| 46 | num_attention_heads=config.get("num_attention_heads", 32), |
| 47 | attention_head_dim=config.get("attention_head_dim", 128), |
| 48 | in_channels=config.get("in_channels", 128), |
| 49 | out_channels=config.get("out_channels", 128), |
| 50 | num_layers=config.get("num_layers", 48), |
| 51 | cross_attention_dim=config.get("cross_attention_dim", 4096), |
| 52 | norm_eps=config.get("norm_eps", 1e-06), |
| 53 | attention_type=AttentionFunction(config.get("attention_type", "default")), |
| 54 | positional_embedding_theta=config.get("positional_embedding_theta", 10000.0), |
| 55 | positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]), |
| 56 | timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000), |
| 57 | use_middle_indices_grid=config.get("use_middle_indices_grid", True), |
| 58 | audio_num_attention_heads=config.get("audio_num_attention_heads", 32), |
| 59 | audio_attention_head_dim=config.get("audio_attention_head_dim", 64), |
| 60 | audio_in_channels=config.get("audio_in_channels", 128), |
| 61 | audio_out_channels=config.get("audio_out_channels", 128), |
| 62 | audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048), |
| 63 | audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]), |
| 64 | av_ca_timestep_scale_multiplier=config.get("av_ca_timestep_scale_multiplier", 1), |
| 65 | rope_type=LTXRopeType(config.get("rope_type", "interleaved")), |
| 66 | double_precision_rope=config.get("frequencies_precision", False) == "float64", |
| 67 | apply_gated_attention=config.get("apply_gated_attention", False), |
| 68 | caption_projection=caption_projection, |
| 69 | audio_caption_projection=audio_caption_projection, |
| 70 | cross_attention_adaln=config.get("cross_attention_adaln", False), |
| 71 | ) |
| 72 | |
| 73 | |
| 74 | class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]): |
| 75 | """ |
| 76 | Configurator for LTX video only model. |
| 77 | Used to create an LTX video only model from a configuration dictionary. |
| 78 | """ |
| 79 | |
| 80 | @classmethod |
| 81 | def from_config(cls: type[LTXModel], config: dict) -> LTXModel: |
| 82 | # Build caption projection for 19B model (projection handled in transformer). |
| 83 | caption_projection, _ = _build_caption_projections(config, is_av=False) |
| 84 | |
| 85 | config = config.get("transformer", {}) |
| 86 | |
| 87 | check_config_value(config, "dropout", 0.0) |
| 88 | check_config_value(config, "attention_bias", True) |
| 89 | check_config_value(config, "num_vector_embeds", None) |
| 90 | check_config_value(config, "activation_fn", "gelu-approximate") |
| 91 | check_config_value(config, "num_embeds_ada_norm", 1000) |
| 92 | check_config_value(config, "use_linear_projection", False) |
| 93 | check_config_value(config, "only_cross_attention", False) |
| 94 | check_config_value(config, "cross_attention_norm", True) |
| 95 | check_config_value(config, "double_self_attention", False) |
| 96 | check_config_value(config, "upcast_attention", False) |
| 97 | check_config_value(config, "standardization_norm", "rms_norm") |
| 98 | check_config_value(config, "norm_elementwise_affine", False) |
| 99 | check_config_value(config, "qk_norm", "rms_norm") |
| 100 | check_config_value(config, "positional_embedding_type", "rope") |
| 101 | check_config_value(config, "use_middle_indices_grid", True) |
| 102 | |
| 103 | return LTXModel( |
| 104 | model_type=LTXModelType.VideoOnly, |
| 105 | num_attention_heads=config.get("num_attention_heads", 32), |
| 106 | attention_head_dim=config.get("attention_head_dim", 128), |
| 107 | in_channels=config.get("in_channels", 128), |
| 108 | out_channels=config.get("out_channels", 128), |
| 109 | num_layers=config.get("num_layers", 48), |
| 110 | cross_attention_dim=config.get("cross_attention_dim", 4096), |
| 111 | norm_eps=config.get("norm_eps", 1e-06), |
| 112 | attention_type=AttentionFunction(config.get("attention_type", "default")), |
| 113 | positional_embedding_theta=config.get("positional_embedding_theta", 10000.0), |
| 114 | positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]), |
| 115 | timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000), |
| 116 | use_middle_indices_grid=config.get("use_middle_indices_grid", True), |
| 117 | rope_type=LTXRopeType(config.get("rope_type", "interleaved")), |
| 118 | double_precision_rope=config.get("frequencies_precision", False) == "float64", |
| 119 | apply_gated_attention=config.get("apply_gated_attention", False), |
| 120 | caption_projection=caption_projection, |
| 121 | cross_attention_adaln=config.get("cross_attention_adaln", False), |
| 122 | ) |
| 123 | |
| 124 | |
| 125 | def _build_caption_projections( |
| 126 | config: dict, |
| 127 | is_av: bool, |
| 128 | ) -> tuple[torch.nn.Module | None, torch.nn.Module | None]: |
| 129 | """Build caption projections for the transformer when projection is NOT in the text encoder. |
| 130 | 19B models: projection is in the transformer (caption_proj_before_connector=False). |
| 131 | 22B models: projection is in the text encoder, so no projections are created here. |
| 132 | Args: |
| 133 | config: Full model config dict (must contain "transformer" key). |
| 134 | is_av: Whether this is an audio-video model. When False, audio projection is skipped. |
| 135 | Returns: |
| 136 | Tuple of (video_caption_projection, audio_caption_projection), both None for 22B models. |
| 137 | """ |
| 138 | transformer_config = config.get("transformer", {}) |
| 139 | if transformer_config.get("caption_proj_before_connector", False): |
| 140 | return None, None |
| 141 | |
| 142 | with torch.device("meta"): |
| 143 | caption_projection = create_caption_projection(transformer_config) |
| 144 | audio_caption_projection = create_caption_projection(transformer_config, audio=True) if is_av else None |
| 145 | return caption_projection, audio_caption_projection |
| 146 | |
| 147 | |
| 148 | LTXV_MODEL_COMFY_RENAMING_MAP = ( |
| 149 | SDOps("LTXV_MODEL_COMFY_PREFIX_MAP") |
| 150 | .with_matching(prefix="model.diffusion_model.") |
| 151 | .with_replacement("model.diffusion_model.", "") |
| 152 | ) |
| 153 |