Support low_cpu_mem_usage=True for the text encoder of Wan2.1 (#146)

---------

Co-authored-by: bubbliiiing <3323290568@qq.com>

scripts/wan2.1/train.py

@@ -842,7 +842,9 @@ def main():
         text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
             additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-        ).to(weight_dtype)
+            low_cpu_mem_usage=True,
+            torch_dtype=weight_dtype,
+        )
         text_encoder = text_encoder.eval()
         # Get Vae
         vae = AutoencoderKLWan.from_pretrained(

scripts/wan2.1/train_lora.py

@@ -841,7 +841,9 @@ def main():
         text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
             additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-        ).to(weight_dtype)
+            low_cpu_mem_usage=True,
+            torch_dtype=weight_dtype,
+        )
         # Get Vae
         vae = AutoencoderKLWan.from_pretrained(
             os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')),

scripts/wan2.1/train_reward_lora.py

@@ -864,7 +864,9 @@ def main():
         text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
             additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-        ).to(weight_dtype)
+            low_cpu_mem_usage=True,
+            torch_dtype=weight_dtype,
+        )
         text_encoder = text_encoder.eval()
         # Get Vae
         vae = AutoencoderKLWan.from_pretrained(

scripts/wan2.1_fun/train.py

@@ -842,7 +842,9 @@ def main():
         text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
             additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-        ).to(weight_dtype)
+            low_cpu_mem_usage=True,
+            torch_dtype=weight_dtype,
+        )
         text_encoder = text_encoder.eval()
         # Get Vae
         vae = AutoencoderKLWan.from_pretrained(

scripts/wan2.1_fun/train_control.py

@@ -799,7 +799,9 @@ def main():
         text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
             additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-        ).to(weight_dtype)
+            low_cpu_mem_usage=True,
+            torch_dtype=weight_dtype,
+        )
         text_encoder = text_encoder.eval()
         # Get Vae
         vae = AutoencoderKLWan.from_pretrained(

scripts/wan2.1_fun/train_lora.py

@@ -840,7 +840,9 @@ def main():
         text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
             additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-        ).to(weight_dtype)
+            low_cpu_mem_usage=True,
+            torch_dtype=weight_dtype,
+        )
         # Get Vae
         vae = AutoencoderKLWan.from_pretrained(
             os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')),

scripts/wan2.1_fun/train_reward_lora.py

@@ -861,7 +861,9 @@ def main():
         text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
             additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-        ).to(weight_dtype)
+            low_cpu_mem_usage=True,
+            torch_dtype=weight_dtype,
+        )
         text_encoder = text_encoder.eval()
         # Get Vae
         vae = AutoencoderKLWan.from_pretrained(