text encoder used from hf

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -24,6 +24,8 @@
 from maxdiffusion.models.embeddings_flax import NNXPixArtAlphaCombinedTimestepSizeEmbeddings, NNXPixArtAlphaTextProjection
 from maxdiffusion.models.gradient_checkpoint import GradientCheckpointType
 from maxdiffusion.common_types import BlockSizes
+from maxdiffusion.configuration_utils import ConfigMixin, register_to_config
+from maxdiffusion.configuration_utils import ConfigMixin, register_to_config
 
 
 class LTX2AdaLayerNormSingle(nnx.Module):
@@ -483,8 +485,9 @@ def __call__(
     return hidden_states, audio_hidden_states
 
 
-class LTX2VideoTransformer3DModel(nnx.Module):
+class LTX2VideoTransformer3DModel(nnx.Module, ConfigMixin):
 
+  @register_to_config
   def __init__(
       self,
       rngs: nnx.Rngs,