diff --git a/comfyui/wan2_1/nodes.py b/comfyui/wan2_1/nodes.py index 5650552..f550356 100644 --- a/comfyui/wan2_1/nodes.py +++ b/comfyui/wan2_1/nodes.py @@ -164,7 +164,9 @@ class LoadWanModel: 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 index c98d1a9..27935e0 100644 --- a/comfyui/wan2_1_fun/nodes.py +++ b/comfyui/wan2_1_fun/nodes.py @@ -163,7 +163,9 @@ class LoadWanFunModel: 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 index 733dffc..3f8aa00 100755 --- a/examples/wan2.1/predict_i2v.py +++ b/examples/wan2.1/predict_i2v.py @@ -135,7 +135,9 @@ tokenizer = AutoTokenizer.from_pretrained( 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 index c8ff4f0..f641c21 100755 --- a/examples/wan2.1/predict_t2v.py +++ b/examples/wan2.1/predict_t2v.py @@ -129,7 +129,9 @@ tokenizer = AutoTokenizer.from_pretrained( 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 index 207d412..d9ab88e 100644 --- a/examples/wan2.1_fun/predict_i2v.py +++ b/examples/wan2.1_fun/predict_i2v.py @@ -135,7 +135,9 @@ tokenizer = AutoTokenizer.from_pretrained( 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 index fbe529c..cee8857 100755 --- a/examples/wan2.1_fun/predict_t2v.py +++ b/examples/wan2.1_fun/predict_t2v.py @@ -130,7 +130,9 @@ tokenizer = AutoTokenizer.from_pretrained( 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 index ada950c..9015331 100644 --- a/examples/wan2.1_fun/predict_v2v_control.py +++ b/examples/wan2.1_fun/predict_v2v_control.py @@ -142,7 +142,9 @@ tokenizer = AutoTokenizer.from_pretrained( 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 index f5f2c03..2b69084 100755 --- a/scripts/wan2.1/train.py +++ b/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( diff --git a/scripts/wan2.1/train_lora.py b/scripts/wan2.1/train_lora.py index 28deeda..c32038d 100755 --- a/scripts/wan2.1/train_lora.py +++ b/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')), diff --git a/scripts/wan2.1/train_reward_lora.py b/scripts/wan2.1/train_reward_lora.py index df780c9..a00b8ad 100755 --- a/scripts/wan2.1/train_reward_lora.py +++ b/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( diff --git a/scripts/wan2.1_fun/train.py b/scripts/wan2.1_fun/train.py index d5ec188..a28c929 100755 --- a/scripts/wan2.1_fun/train.py +++ b/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( diff --git a/scripts/wan2.1_fun/train_control.py b/scripts/wan2.1_fun/train_control.py index a929d16..95e7615 100755 --- a/scripts/wan2.1_fun/train_control.py +++ b/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( diff --git a/scripts/wan2.1_fun/train_lora.py b/scripts/wan2.1_fun/train_lora.py index c1e0786..4bd5b08 100755 --- a/scripts/wan2.1_fun/train_lora.py +++ b/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')), diff --git a/scripts/wan2.1_fun/train_reward_lora.py b/scripts/wan2.1_fun/train_reward_lora.py index 1ff4eb9..e24d9e4 100755 --- a/scripts/wan2.1_fun/train_reward_lora.py +++ b/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( diff --git a/videox_fun/api/api_multi_nodes.py b/videox_fun/api/api_multi_nodes.py index 005d8eb..5dbcd52 100755 --- 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 index 7ec5349..34a0323 100644 --- a/videox_fun/models/wan_text_encoder.py +++ b/videox_fun/models/wan_text_encoder.py @@ -304,21 +304,62 @@ class WanT5EncoderModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): 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 + + if low_cpu_mem_usage: + try: + import re - 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 \ No newline at end of file + 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." + ) \ No newline at end of file diff --git a/videox_fun/ui/wan_fun_ui.py b/videox_fun/ui/wan_fun_ui.py index 82d80f9..6f7b629 100644 --- a/videox_fun/ui/wan_fun_ui.py +++ b/videox_fun/ui/wan_fun_ui.py @@ -65,7 +65,9 @@ class Wan_Fun_Controller(Fun_Controller): 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 index f8204ee..621726e 100644 --- a/videox_fun/ui/wan_ui.py +++ b/videox_fun/ui/wan_ui.py @@ -65,7 +65,9 @@ class Wan_Controller(Fun_Controller): 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: