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

hkunzhe · bubbliiiing · web-flow · commit 09e7bd2bd01f · 2025-04-01T11:35:42.000+08:00
---------

Co-authored-by: bubbliiiing &lt;3323290568@qq.com&gt;
diff --git a/comfyui/wan2_1/nodes.py b/comfyui/wan2_1/nodes.py
@@ -164,7 +164,9 @@ def loadmodel(self, GPU_memory_mode, model, precision, config):
         text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(model_name, 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,
+        )
         pbar.update(1) 
 
         if transformer.config.in_channels != vae.config.latent_channels:
diff --git a/comfyui/wan2_1_fun/nodes.py b/comfyui/wan2_1_fun/nodes.py
@@ -163,7 +163,9 @@ def loadmodel(self, GPU_memory_mode, model_type, model, precision, config):
         text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(model_name, 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,
+        )
         pbar.update(1) 
 
         if transformer.config.in_channels != vae.config.latent_channels:
diff --git a/examples/wan2.1/predict_i2v.py b/examples/wan2.1/predict_i2v.py
@@ -135,7 +135,9 @@
 text_encoder = WanT5EncoderModel.from_pretrained(
     os.path.join(model_name, 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 Clip Image Encoder
diff --git a/examples/wan2.1/predict_t2v.py b/examples/wan2.1/predict_t2v.py
@@ -129,7 +129,9 @@
 text_encoder = WanT5EncoderModel.from_pretrained(
     os.path.join(model_name, 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 Scheduler
 Choosen_Scheduler = scheduler_dict = {
diff --git a/examples/wan2.1_fun/predict_i2v.py b/examples/wan2.1_fun/predict_i2v.py
@@ -135,7 +135,9 @@
 text_encoder = WanT5EncoderModel.from_pretrained(
     os.path.join(model_name, 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 Clip Image Encoder
diff --git a/examples/wan2.1_fun/predict_t2v.py b/examples/wan2.1_fun/predict_t2v.py
@@ -130,7 +130,9 @@
 text_encoder = WanT5EncoderModel.from_pretrained(
     os.path.join(model_name, 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()
 
 if transformer.config.in_channels != vae.config.latent_channels:
diff --git a/examples/wan2.1_fun/predict_v2v_control.py b/examples/wan2.1_fun/predict_v2v_control.py
@@ -142,7 +142,9 @@
 text_encoder = WanT5EncoderModel.from_pretrained(
     os.path.join(model_name, 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 Clip Image Encoder
diff --git a/scripts/wan2.1/train.py b/scripts/wan2.1/train.py
@@ -842,7 +842,9 @@ def deepspeed_zero_init_disabled_context_manager():
         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(
diff --git a/scripts/wan2.1/train_lora.py b/scripts/wan2.1/train_lora.py
@@ -841,7 +841,9 @@ def deepspeed_zero_init_disabled_context_manager():
         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')),
diff --git a/scripts/wan2.1/train_reward_lora.py b/scripts/wan2.1/train_reward_lora.py
@@ -864,7 +864,9 @@ def deepspeed_zero_init_disabled_context_manager():
         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(
diff --git a/scripts/wan2.1_fun/train.py b/scripts/wan2.1_fun/train.py
@@ -842,7 +842,9 @@ def deepspeed_zero_init_disabled_context_manager():
         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(
diff --git a/scripts/wan2.1_fun/train_control.py b/scripts/wan2.1_fun/train_control.py
@@ -799,7 +799,9 @@ def deepspeed_zero_init_disabled_context_manager():
         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(
diff --git a/scripts/wan2.1_fun/train_lora.py b/scripts/wan2.1_fun/train_lora.py
@@ -840,7 +840,9 @@ def deepspeed_zero_init_disabled_context_manager():
         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')),
diff --git a/scripts/wan2.1_fun/train_reward_lora.py b/scripts/wan2.1_fun/train_reward_lora.py
@@ -861,7 +861,9 @@ def deepspeed_zero_init_disabled_context_manager():
         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(
diff --git a/videox_fun/api/api_multi_nodes.py b/videox_fun/api/api_multi_nodes.py
@@ -1,4 +1,4 @@
-
+# This file is modified from https://github.com/xdit-project/xDiT/blob/main/entrypoints/launch.py
 import base64
 import gc
 import os
diff --git a/videox_fun/models/wan_text_encoder.py b/videox_fun/models/wan_text_encoder.py
@@ -304,21 +304,62 @@ def forward(
         return (x, )
     
     @classmethod
-    def from_pretrained(cls, pretrained_model_path, additional_kwargs={}):
+    def from_pretrained(cls, pretrained_model_path, additional_kwargs={}, low_cpu_mem_usage=False, torch_dtype=torch.bfloat16):
         def filter_kwargs(cls, kwargs):
             import inspect
             sig = inspect.signature(cls.__init__)
             valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
             filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
             return filtered_kwargs
-
-        model = cls(**filter_kwargs(cls, additional_kwargs))
-        if pretrained_model_path.endswith(".safetensors"):
-            from safetensors.torch import load_file, safe_open
-            state_dict = load_file(pretrained_model_path)
-        else:
-            state_dict = torch.load(pretrained_model_path, map_location="cpu")
-        m, u = model.load_state_dict(state_dict, strict=False)
-        print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
-        print(m, u)
-        return model
+    
+        if low_cpu_mem_usage:
+            try:
+                import re
+
+                from diffusers.models.modeling_utils import \
+                    load_model_dict_into_meta
+                from diffusers.utils import is_accelerate_available
+                if is_accelerate_available():
+                    import accelerate
+                
+                # Instantiate model with empty weights
+                with accelerate.init_empty_weights():
+                    model = cls(**filter_kwargs(cls, additional_kwargs))
+
+                param_device = "cpu"
+                if pretrained_model_path.endswith(".safetensors"):
+                    from safetensors.torch import load_file
+                    state_dict = load_file(pretrained_model_path)
+                else:
+                    state_dict = torch.load(pretrained_model_path, map_location="cpu")
+                # move the params from meta device to cpu
+                missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
+                if len(missing_keys) > 0:
+                    raise ValueError(
+                        f"Cannot load {cls} from {pretrained_model_path} because the following keys are"
+                        f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"
+                        " `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"
+                        " those weights or else make sure your checkpoint file is correct."
+                    )
+
+                unexpected_keys = load_model_dict_into_meta(
+                    model,
+                    state_dict,
+                    device=param_device,
+                    dtype=torch_dtype,
+                    model_name_or_path=pretrained_model_path,
+                )
+
+                if cls._keys_to_ignore_on_load_unexpected is not None:
+                    for pat in cls._keys_to_ignore_on_load_unexpected:
+                        unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
+
+                if len(unexpected_keys) > 0:
+                    print(
+                        f"Some weights of the model checkpoint were not used when initializing {cls.__name__}: \n {[', '.join(unexpected_keys)]}"
+                    )
+                return model
+            except Exception as e:
+                print(
+                    f"The low_cpu_mem_usage mode is not work because {e}. Use low_cpu_mem_usage=False instead."
+                )
diff --git a/videox_fun/ui/wan_fun_ui.py b/videox_fun/ui/wan_fun_ui.py
@@ -65,7 +65,9 @@ def update_diffusion_transformer(self, diffusion_transformer_dropdown):
         self.text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(diffusion_transformer_dropdown, self.config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
             additional_kwargs=OmegaConf.to_container(self.config['text_encoder_kwargs']),
-        ).to(self.weight_dtype)
+            low_cpu_mem_usage=True,
+            torch_dtype=self.weight_dtype,
+        )
         self.text_encoder = self.text_encoder.eval()
 
         if self.transformer.config.in_channels != self.vae.config.latent_channels:
diff --git a/videox_fun/ui/wan_ui.py b/videox_fun/ui/wan_ui.py
@@ -65,7 +65,9 @@ def update_diffusion_transformer(self, diffusion_transformer_dropdown):
         self.text_encoder = WanT5EncoderModel.from_pretrained(
             os.path.join(diffusion_transformer_dropdown, self.config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
             additional_kwargs=OmegaConf.to_container(self.config['text_encoder_kwargs']),
-        ).to(self.weight_dtype)
+            low_cpu_mem_usage=True,
+            torch_dtype=self.weight_dtype,
+        )
         self.text_encoder = self.text_encoder.eval()
 
         if self.transformer.config.in_channels != self.vae.config.latent_channels:
