| 1 | from ltx_core.loader.sd_ops import SDOps |
| 2 | from ltx_core.model.model_protocol import ModelConfigurator |
| 3 | from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, PaddingModeType |
| 4 | from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder |
| 5 | |
| 6 | |
| 7 | class VideoEncoderConfigurator(ModelConfigurator[VideoEncoder]): |
| 8 | """Configurator for creating a video VAE Encoder from a configuration dictionary.""" |
| 9 | |
| 10 | @classmethod |
| 11 | def from_config(cls: type[VideoEncoder], config: dict) -> VideoEncoder: |
| 12 | config = config.get("vae", {}) |
| 13 | convolution_dimensions = config.get("dims", 3) |
| 14 | in_channels = config.get("in_channels", 3) |
| 15 | latent_channels = config.get("latent_channels", 128) |
| 16 | spatial_padding_mode = PaddingModeType(config.get("spatial_padding_mode", "zeros")) |
| 17 | encoder_blocks = config.get("encoder_blocks", []) |
| 18 | patch_size = config.get("patch_size", 4) |
| 19 | norm_layer_str = config.get("norm_layer", "pixel_norm") |
| 20 | latent_log_var_str = config.get("latent_log_var", "uniform") |
| 21 | |
| 22 | return VideoEncoder( |
| 23 | convolution_dimensions=convolution_dimensions, |
| 24 | in_channels=in_channels, |
| 25 | out_channels=latent_channels, |
| 26 | encoder_blocks=encoder_blocks, |
| 27 | patch_size=patch_size, |
| 28 | norm_layer=NormLayerType(norm_layer_str), |
| 29 | latent_log_var=LogVarianceType(latent_log_var_str), |
| 30 | encoder_spatial_padding_mode=spatial_padding_mode, |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | class VideoDecoderConfigurator(ModelConfigurator[VideoDecoder]): |
| 35 | """Configurator for creating a video VAE Decoder from a configuration dictionary.""" |
| 36 | |
| 37 | @classmethod |
| 38 | def from_config(cls: type[VideoDecoder], config: dict) -> VideoDecoder: |
| 39 | config = config.get("vae", {}) |
| 40 | convolution_dimensions = config.get("dims", 3) |
| 41 | latent_channels = config.get("latent_channels", 128) |
| 42 | spatial_padding_mode = PaddingModeType(config.get("spatial_padding_mode", "reflect")) |
| 43 | out_channels = config.get("out_channels", 3) |
| 44 | decoder_blocks = config.get("decoder_blocks", []) |
| 45 | patch_size = config.get("patch_size", 4) |
| 46 | norm_layer_str = config.get("norm_layer", "pixel_norm") |
| 47 | causal = config.get("causal_decoder", False) |
| 48 | timestep_conditioning = config.get("timestep_conditioning", True) |
| 49 | base_channels = config.get("decoder_base_channels", 128) |
| 50 | |
| 51 | return VideoDecoder( |
| 52 | convolution_dimensions=convolution_dimensions, |
| 53 | in_channels=latent_channels, |
| 54 | out_channels=out_channels, |
| 55 | decoder_blocks=decoder_blocks, |
| 56 | patch_size=patch_size, |
| 57 | norm_layer=NormLayerType(norm_layer_str), |
| 58 | causal=causal, |
| 59 | timestep_conditioning=timestep_conditioning, |
| 60 | decoder_spatial_padding_mode=spatial_padding_mode, |
| 61 | base_channels=base_channels, |
| 62 | ) |
| 63 | |
| 64 | |
| 65 | VAE_DECODER_COMFY_KEYS_FILTER = ( |
| 66 | SDOps("VAE_DECODER_COMFY_KEYS_FILTER") |
| 67 | .with_matching(prefix="vae.decoder.") |
| 68 | .with_matching(prefix="vae.per_channel_statistics.") |
| 69 | .with_replacement("vae.decoder.", "") |
| 70 | .with_replacement("vae.per_channel_statistics.", "per_channel_statistics.") |
| 71 | ) |
| 72 | |
| 73 | VAE_ENCODER_COMFY_KEYS_FILTER = ( |
| 74 | SDOps("VAE_ENCODER_COMFY_KEYS_FILTER") |
| 75 | .with_matching(prefix="vae.encoder.") |
| 76 | .with_matching(prefix="vae.per_channel_statistics.") |
| 77 | .with_replacement("vae.encoder.", "") |
| 78 | .with_replacement("vae.per_channel_statistics.", "per_channel_statistics.") |
| 79 | ) |
| 80 |