From acffc2d9fcefd23244cd6e897b9c7f8db04c093d Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 9 Jan 2025 12:54:57 +0530 Subject: [PATCH 1/5] parse subfolder when needed. --- finetrainers/cogvideox/cogvideox_lora.py | 3 ++- finetrainers/hunyuan_video/hunyuan_video_lora.py | 3 ++- finetrainers/ltx_video/ltx_video_lora.py | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/finetrainers/cogvideox/cogvideox_lora.py b/finetrainers/cogvideox/cogvideox_lora.py index c3b754a..9d3a77c 100644 --- a/finetrainers/cogvideox/cogvideox_lora.py +++ b/finetrainers/cogvideox/cogvideox_lora.py @@ -42,8 +42,9 @@ def load_diffusion_models( cache_dir: Optional[str] = None, **kwargs, ): + subfolder = kwargs.get("subfolder", None) transformer = CogVideoXTransformer3DModel.from_pretrained( - model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir + model_id, subfolder=subfolder, torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir ) scheduler = CogVideoXDDIMScheduler.from_pretrained(model_id, subfolder="scheduler") return {"transformer": transformer, "scheduler": scheduler} diff --git a/finetrainers/hunyuan_video/hunyuan_video_lora.py b/finetrainers/hunyuan_video/hunyuan_video_lora.py index 9bfea53..450101d 100644 --- a/finetrainers/hunyuan_video/hunyuan_video_lora.py +++ b/finetrainers/hunyuan_video/hunyuan_video_lora.py @@ -63,8 +63,9 @@ def load_diffusion_models( cache_dir: Optional[str] = None, **kwargs, ) -> Dict[str, Union[nn.Module, FlowMatchEulerDiscreteScheduler]]: + subfolder = kwargs.get("subfolder", None) transformer = HunyuanVideoTransformer3DModel.from_pretrained( - model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir + model_id, subfolder=subfolder, torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir ) scheduler = FlowMatchEulerDiscreteScheduler(shift=shift) return {"transformer": transformer, "scheduler": scheduler} diff --git a/finetrainers/ltx_video/ltx_video_lora.py b/finetrainers/ltx_video/ltx_video_lora.py index 0e1af9b..7a80543 100644 --- a/finetrainers/ltx_video/ltx_video_lora.py +++ b/finetrainers/ltx_video/ltx_video_lora.py @@ -45,8 +45,9 @@ def load_diffusion_models( cache_dir: Optional[str] = None, **kwargs, ) -> Dict[str, nn.Module]: + subfolder = kwargs.get("subfolder", None) transformer = LTXVideoTransformer3DModel.from_pretrained( - model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir + model_id, subfolder=subfolder, torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir ) scheduler = FlowMatchEulerDiscreteScheduler() return {"transformer": transformer, "scheduler": scheduler} From 8188f8a11e631ddf59335a76d714df76bf0b982a Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 9 Jan 2025 13:03:14 +0530 Subject: [PATCH 2/5] tackle saving and loading hooks. --- finetrainers/trainer.py | 94 +++++++++++++++++++++++------------------ 1 file changed, 54 insertions(+), 40 deletions(-) diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index 81c8264..125e70d 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -406,13 +406,16 @@ class Trainer: if self.args.gradient_checkpointing: self.transformer.enable_gradient_checkpointing() - transformer_lora_config = LoraConfig( - r=self.args.rank, - lora_alpha=self.args.lora_alpha, - init_lora_weights=True, - target_modules=self.args.target_modules, - ) - self.transformer.add_adapter(transformer_lora_config) + if self.args.training_type == "lora": + transformer_lora_config = LoraConfig( + r=self.args.rank, + lora_alpha=self.args.lora_alpha, + init_lora_weights=True, + target_modules=self.args.target_modules, + ) + self.transformer.add_adapter(transformer_lora_config) + else: + transformer_lora_config = None # Enable TF32 for faster training on Ampere GPUs: https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices if self.args.allow_tf32 and torch.cuda.is_available(): @@ -432,7 +435,8 @@ class Trainer: type(unwrap_model(self.state.accelerator, self.transformer)), ): model = unwrap_model(self.state.accelerator, model) - transformer_lora_layers_to_save = get_peft_model_state_dict(model) + if self.args.training_type == "lora": + transformer_lora_layers_to_save = get_peft_model_state_dict(model) else: raise ValueError(f"Unexpected save model: {model.__class__}") @@ -440,10 +444,14 @@ class Trainer: if weights: weights.pop() - self.model_config["pipeline_cls"].save_lora_weights( - output_dir, - transformer_lora_layers=transformer_lora_layers_to_save, - ) + # TODO: refactor later if needed. But for now, this is just a few LoC. + if self.args.training_type == "lora": + self.model_config["pipeline_cls"].save_lora_weights( + output_dir, + transformer_lora_layers=transformer_lora_layers_to_save, + ) + else: + model.save_pretrained(os.path.join(output_dir, "transformer")) def load_model_hook(models, input_dir): if not self.state.accelerator.distributed_type == DistributedType.DEEPSPEED: @@ -459,33 +467,39 @@ class Trainer: f"Unexpected save model: {unwrap_model(self.state.accelerator, model).__class__}" ) else: - transformer_ = unwrap_model(self.state.accelerator, self.transformer).__class__.from_pretrained( - self.args.pretrained_model_name_or_path, subfolder="transformer" - ) - transformer_.add_adapter(transformer_lora_config) + transformer_cls_ = unwrap_model(self.state.accelerator, self.transformer).__class__ - lora_state_dict = self.model_config["pipeline_cls"].lora_state_dict(input_dir) - transformer_state_dict = { - f'{k.replace("transformer.", "")}': v - for k, v in lora_state_dict.items() - if k.startswith("transformer.") - } - incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default") - if incompatible_keys is not None: - # check only for unexpected keys - unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) - if unexpected_keys: - logger.warning( - f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " - f" {unexpected_keys}. " + if self.args.training_type == "lora": + transformer_ = transformer_cls_.__class__.from_pretrained( + self.args.pretrained_model_name_or_path, subfolder="transformer" ) + transformer_.add_adapter(transformer_lora_config) + lora_state_dict = self.model_config["pipeline_cls"].lora_state_dict(input_dir) + transformer_state_dict = { + f'{k.replace("transformer.", "")}': v + for k, v in lora_state_dict.items() + if k.startswith("transformer.") + } + incompatible_keys = set_peft_model_state_dict( + transformer_, transformer_state_dict, adapter_name="default" + ) + if incompatible_keys is not None: + # check only for unexpected keys + unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) + if unexpected_keys: + logger.warning( + f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " + f" {unexpected_keys}. " + ) - # Make sure the trainable params are in float32. This is again needed since the base models - # are in `weight_dtype`. More details: - # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804 - if self.args.mixed_precision == "fp16": - # only upcast trainable parameters (LoRA) into fp32 - cast_training_params([transformer_]) + # Make sure the trainable params are in float32. This is again needed since the base models + # are in `weight_dtype`. More details: + # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804 + if self.args.mixed_precision == "fp16": + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params([transformer_]) + else: + transformer_ = transformer_cls_.from_pretrained(os.path.join(input_dir, "transformer")) self.state.accelerator.register_save_state_pre_hook(save_model_hook) self.state.accelerator.register_load_state_pre_hook(load_model_hook) @@ -497,7 +511,7 @@ class Trainer: self.state.train_steps = self.args.train_steps # Make sure the trainable params are in float32 - if self.args.mixed_precision == "fp16": + if self.args.mixed_precision == "fp16" and self.args.training_type == "lora": # only upcast trainable parameters (LoRA) into fp32 cast_training_params([self.transformer], dtype=torch.float32) @@ -510,13 +524,13 @@ class Trainer: * self.state.accelerator.num_processes ) - transformer_lora_parameters = list(filter(lambda p: p.requires_grad, self.transformer.parameters())) + transformer_trainable_parameters = list(filter(lambda p: p.requires_grad, self.transformer.parameters())) transformer_parameters_with_lr = { - "params": transformer_lora_parameters, + "params": transformer_trainable_parameters, "lr": self.state.learning_rate, } params_to_optimize = [transformer_parameters_with_lr] - self.state.num_trainable_parameters = sum(p.numel() for p in transformer_lora_parameters) + self.state.num_trainable_parameters = sum(p.numel() for p in transformer_trainable_parameters) use_deepspeed_opt = ( self.state.accelerator.state.deepspeed_plugin is not None From 5183405e61dfa2cdf1d56abbd46cd1d588b672f7 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 9 Jan 2025 13:17:15 +0530 Subject: [PATCH 3/5] tackle validation. --- finetrainers/trainer.py | 73 +++++++++++++++++++++++++---------------- 1 file changed, 44 insertions(+), 29 deletions(-) diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index 125e70d..502d2b4 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -21,6 +21,7 @@ from accelerate.utils import ( gather_object, set_seed, ) +from diffusers import DiffusionPipeline from diffusers.configuration_utils import FrozenDict from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.optimization import get_scheduler @@ -916,35 +917,7 @@ class Trainer: memory_statistics = get_memory_statistics() logger.info(f"Memory before validation start: {json.dumps(memory_statistics, indent=4)}") - if not final_validation: - pipeline = self.model_config["initialize_pipeline"]( - model_id=self.args.pretrained_model_name_or_path, - tokenizer=self.tokenizer, - text_encoder=self.text_encoder, - tokenizer_2=self.tokenizer_2, - text_encoder_2=self.text_encoder_2, - transformer=unwrap_model(accelerator, self.transformer), - vae=self.vae, - device=accelerator.device, - revision=self.args.revision, - cache_dir=self.args.cache_dir, - enable_slicing=self.args.enable_slicing, - enable_tiling=self.args.enable_tiling, - enable_model_cpu_offload=self.args.enable_model_cpu_offload, - ) - else: - # `torch_dtype` is manually set within `initialize_pipeline()`. - self._delete_components() - pipeline = self.model_config["initialize_pipeline"]( - model_id=self.args.pretrained_model_name_or_path, - device=accelerator.device, - revision=self.args.revision, - cache_dir=self.args.cache_dir, - enable_slicing=self.args.enable_slicing, - enable_tiling=self.args.enable_tiling, - enable_model_cpu_offload=self.args.enable_model_cpu_offload, - ) - pipeline.load_lora_weights(self.args.output_dir) + pipeline = self._get_and_prepare_pipeline_for_validation(final_validation=final_validation) all_processes_artifacts = [] prompts_to_filenames = {} @@ -1150,3 +1123,45 @@ class Trainer: elif self.state.accelerator.mixed_precision == "bf16": weight_dtype = torch.bfloat16 return weight_dtype + + def _get_and_prepare_pipeline_for_validation(self, final_validation: bool = False) -> DiffusionPipeline: + accelerator = self.state.accelerator + if not final_validation: + pipeline = self.model_config["initialize_pipeline"]( + model_id=self.args.pretrained_model_name_or_path, + tokenizer=self.tokenizer, + text_encoder=self.text_encoder, + tokenizer_2=self.tokenizer_2, + text_encoder_2=self.text_encoder_2, + transformer=unwrap_model(accelerator, self.transformer), + vae=self.vae, + device=accelerator.device, + revision=self.args.revision, + cache_dir=self.args.cache_dir, + enable_slicing=self.args.enable_slicing, + enable_tiling=self.args.enable_tiling, + enable_model_cpu_offload=self.args.enable_model_cpu_offload, + ) + else: + # `torch_dtype` is manually set within `initialize_pipeline()`. + self._delete_components() + if self.args.training_type == "lora": + transformer = None + else: + transformer = self.model_config["load_diffusion_models"]( + model_id=self.args.output_dir, subfolder=None + )["transformer"] + pipeline = self.model_config["initialize_pipeline"]( + model_id=self.args.pretrained_model_name_or_path, + transformer=transformer, + device=accelerator.device, + revision=self.args.revision, + cache_dir=self.args.cache_dir, + enable_slicing=self.args.enable_slicing, + enable_tiling=self.args.enable_tiling, + enable_model_cpu_offload=self.args.enable_model_cpu_offload, + ) + if self.args.training_type == "lora": + pipeline.load_lora_weights(self.args.output_dir) + + return pipeline From 162e6cd50dc745656d4685b3159c7446f4fa7dc5 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 9 Jan 2025 13:25:01 +0530 Subject: [PATCH 4/5] fix subfolder bug. --- finetrainers/cogvideox/cogvideox_lora.py | 2 +- finetrainers/hunyuan_video/hunyuan_video_lora.py | 2 +- finetrainers/ltx_video/ltx_video_lora.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/finetrainers/cogvideox/cogvideox_lora.py b/finetrainers/cogvideox/cogvideox_lora.py index 9d3a77c..c1413ae 100644 --- a/finetrainers/cogvideox/cogvideox_lora.py +++ b/finetrainers/cogvideox/cogvideox_lora.py @@ -42,7 +42,7 @@ def load_diffusion_models( cache_dir: Optional[str] = None, **kwargs, ): - subfolder = kwargs.get("subfolder", None) + subfolder = kwargs.get("subfolder", "transformer") transformer = CogVideoXTransformer3DModel.from_pretrained( model_id, subfolder=subfolder, torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir ) diff --git a/finetrainers/hunyuan_video/hunyuan_video_lora.py b/finetrainers/hunyuan_video/hunyuan_video_lora.py index 450101d..c23feb8 100644 --- a/finetrainers/hunyuan_video/hunyuan_video_lora.py +++ b/finetrainers/hunyuan_video/hunyuan_video_lora.py @@ -63,7 +63,7 @@ def load_diffusion_models( cache_dir: Optional[str] = None, **kwargs, ) -> Dict[str, Union[nn.Module, FlowMatchEulerDiscreteScheduler]]: - subfolder = kwargs.get("subfolder", None) + subfolder = kwargs.get("subfolder", "transformer") transformer = HunyuanVideoTransformer3DModel.from_pretrained( model_id, subfolder=subfolder, torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir ) diff --git a/finetrainers/ltx_video/ltx_video_lora.py b/finetrainers/ltx_video/ltx_video_lora.py index 7a80543..a6e8222 100644 --- a/finetrainers/ltx_video/ltx_video_lora.py +++ b/finetrainers/ltx_video/ltx_video_lora.py @@ -45,7 +45,7 @@ def load_diffusion_models( cache_dir: Optional[str] = None, **kwargs, ) -> Dict[str, nn.Module]: - subfolder = kwargs.get("subfolder", None) + subfolder = kwargs.get("subfolder", "transformer") transformer = LTXVideoTransformer3DModel.from_pretrained( model_id, subfolder=subfolder, torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir ) From c0f38892f88158d77fdeb2fc32d5eb3742a71473 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 10 Jan 2025 07:44:36 +0530 Subject: [PATCH 5/5] remove __class__. --- finetrainers/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index 922186d..e413fe8 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -471,7 +471,7 @@ class Trainer: transformer_cls_ = unwrap_model(self.state.accelerator, self.transformer).__class__ if self.args.training_type == "lora": - transformer_ = transformer_cls_.__class__.from_pretrained( + transformer_ = transformer_cls_.from_pretrained( self.args.pretrained_model_name_or_path, subfolder="transformer" ) transformer_.add_adapter(transformer_lora_config)