sampled video looks correct

Author: jzhang38

fastvideo/train/methods/rl/embeddings.py

@@ -34,13 +34,12 @@ def compute_text_embeddings(
         return_tensors="pt",
     )
     text_input_ids = text_inputs.input_ids.to(device)
-
+    attention_mask = text_inputs.attention_mask.to(device)
     with torch.no_grad():
         prompt_embeds = text_encoder(
-            text_input_ids
-        )[0]
-
-    prompt_embeds = prompt_embeds.to(
-        dtype=text_encoder.dtype, device=device
-    )
+            text_input_ids, attention_mask=attention_mask
+        ).last_hidden_state
+    
+    # make padding token 0
+    prompt_embeds[attention_mask == 0] = 0
     return prompt_embeds

fastvideo/train/models/wan/wan_genrl.py

@@ -104,17 +104,20 @@ def init_preprocessors(
 
         # Load text encoder and tokenizer.
         model_path = str(training_config.model_path)
-        self._load_text_encoder(model_path)
+        self._load_text_encoder(
+            model_path, training_config
+        )
 
         # Dummy dataloader for the trainer's outer loop.
         self.dataloader = _InfiniteDummyLoader()
         self.start_step = 0
 
-    def _load_text_encoder(self, model_path: str) -> None:
-        from transformers import (
-            AutoTokenizer,
-            UMT5EncoderModel,
-        )
+    def _load_text_encoder(
+        self,
+        model_path: str,
+        training_config: TrainingConfig,
+    ) -> None:
+        from transformers import AutoTokenizer
 
         logger.info(
             "Loading tokenizer from %s", model_path
@@ -124,17 +127,14 @@ def _load_text_encoder(self, model_path: str) -> None:
         )
 
         logger.info(
-            "Loading T5 text encoder from %s", model_path
+            "Loading text encoder from %s", model_path
         )
-        dtype = self._get_training_dtype()
-        self.text_encoder = UMT5EncoderModel.from_pretrained(
-            model_path,
-            subfolder="text_encoder",
-            torch_dtype=dtype,
+        self.text_encoder = load_module_from_path(
+            model_path=model_path,
+            module_type="text_encoder",
+            training_config=training_config,
         )
-        self.text_encoder.to(self.device)
         self.text_encoder.requires_grad_(False)
-        self.text_encoder.eval()
 
     def on_train_start(self) -> None:
         """Skip negative conditioning (handled by method)."""