Support train TE

hunyuan_train_network.py

@@ -1,3 +1,7 @@
+import warnings
+
+warnings.filterwarnings("ignore")
+
 import argparse
 
 import torch
@@ -103,6 +107,8 @@ def get_text_cond(
         ):
             input_ids1 = batch["input_ids"]
             input_ids2 = batch["input_ids2"]
+            logger.debug("input_ids1", input_ids1.shape)
+            logger.debug("input_ids2", input_ids2.shape)
             with torch.enable_grad():
                 input_ids1 = input_ids1.to(accelerator.device)
                 input_ids2 = input_ids2.to(accelerator.device)
@@ -119,6 +125,8 @@ def get_text_cond(
                         accelerator=accelerator,
                     )
                 )
+                logger.debug("encoder_hidden_states1", encoder_hidden_states1.shape)
+                logger.debug("encoder_hidden_states2", encoder_hidden_states2.shape)
         else:
             raise NotImplementedError
         return encoder_hidden_states1, mask1, encoder_hidden_states2, mask2

train_network.py

@@ -477,7 +477,15 @@ def train(self, args):
 
                 # set top parameter requires_grad = True for gradient checkpointing works
                 if train_text_encoder:
-                    t_enc.text_model.embeddings.requires_grad_(True)
+                    if hasattr(t_enc, "text_model"):
+                        t_enc.text_model.embeddings.requires_grad_(True)
+                    elif hasattr(t_enc, "embeddings"):
+                        # HunYuan Bert(CLIP)
+                        t_enc.embeddings.requires_grad_(True)
+                    elif hasattr(t_enc, "get_token_embedding"):
+                        # Others (mT5 or other encoder, will have custom method to get the correct embedding)
+                        t_enc.get_token_embedding().requires_grad_(True)
+
 
         else:
             unet.eval()