weight loading for tokenizer

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -328,7 +328,7 @@ def load_vae_weights(
 
       flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict)
       flax_key = _tuple_str_to_int(flax_key)
-      max_logging.log(f"Mapped key: {pt_key} -> {flax_key}")
+
 
       if resnet_index is not None:
         str_flax_key = tuple([str(x) for x in flax_key])

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -307,7 +307,7 @@ def load_tokenizer(cls, config: HyperParameters):
   def load_text_encoder(cls, config: HyperParameters):
     max_logging.log("Loading Gemma3 Text Encoder...")
     text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
-        config.pretrained_model_name_or_path,
+        "Lightricks/LTX-2",
         subfolder="text_encoder",
         torch_dtype=torch.bfloat16,
     )