diff --git a/finetrainers/args.py b/finetrainers/args.py index bcd0076..137eab6 100644 --- a/finetrainers/args.py +++ b/finetrainers/args.py @@ -236,6 +236,7 @@ def parse_arguments() -> Args: def validate_args(args: Args): + _validate_training_args(args) _validate_validation_args(args) @@ -455,7 +456,8 @@ def _add_training_arguments(parser: argparse.ArgumentParser) -> None: "--training_type", type=str, default=None, - help="Type of training to perform. Choose between ['lora','sft']", + choices=["lora", "full-finetune"], + help="Type of training to perform. Choose between ['lora', 'full-finetune']", ) parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") parser.add_argument( @@ -874,6 +876,15 @@ def _map_to_args_type(args: Dict[str, Any]) -> Args: return result_args +def _validate_training_args(args: Args): + if args.training_type == "lora": + assert args.rank is not None, "Rank is required for LoRA training" + assert args.lora_alpha is not None, "LoRA alpha is required for LoRA training" + assert ( + args.target_modules is not None and len(args.target_modules) > 0 + ), "Target modules are required for LoRA training" + + def _validate_validation_args(args: Args): assert args.validation_prompts is not None, "Validation prompts are required for validation" if args.validation_images is not None: diff --git a/finetrainers/cogvideox/__init__.py b/finetrainers/cogvideox/__init__.py index 8ca65b7..390479b 100644 --- a/finetrainers/cogvideox/__init__.py +++ b/finetrainers/cogvideox/__init__.py @@ -1 +1,2 @@ -from .cogvideox_lora import COGVIDEOX_T2V_FULL_FINETUNE_CONFIG, COGVIDEOX_T2V_LORA_CONFIG +from .cogvideox_lora import COGVIDEOX_T2V_LORA_CONFIG +from .full_finetune import COGVIDEOX_T2V_FULL_FINETUNE_CONFIG diff --git a/finetrainers/cogvideox/cogvideox_lora.py b/finetrainers/cogvideox/cogvideox_lora.py index cc213f8..36776d0 100644 --- a/finetrainers/cogvideox/cogvideox_lora.py +++ b/finetrainers/cogvideox/cogvideox_lora.py @@ -312,6 +312,7 @@ def _pad_frames(latents: torch.Tensor, patch_size_t: int): return latents +# TODO(aryan): refactor into model specs for better re-use COGVIDEOX_T2V_LORA_CONFIG = { "pipeline_cls": CogVideoXPipeline, "load_condition_models": load_condition_models, @@ -326,18 +327,3 @@ COGVIDEOX_T2V_LORA_CONFIG = { "forward_pass": forward_pass, "validation": validation, } - -COGVIDEOX_T2V_FULL_FINETUNE_CONFIG = { - "pipeline_cls": CogVideoXPipeline, - "load_condition_models": load_condition_models, - "load_latent_models": load_latent_models, - "load_diffusion_models": load_diffusion_models, - "initialize_pipeline": initialize_pipeline, - "prepare_conditions": prepare_conditions, - "prepare_latents": prepare_latents, - "post_latent_preparation": post_latent_preparation, - "collate_fn": collate_fn_t2v, - "calculate_noisy_latents": calculate_noisy_latents, - "forward_pass": forward_pass, - "validation": validation, -} diff --git a/finetrainers/cogvideox/full_finetune.py b/finetrainers/cogvideox/full_finetune.py new file mode 100644 index 0000000..f755981 --- /dev/null +++ b/finetrainers/cogvideox/full_finetune.py @@ -0,0 +1,32 @@ +from diffusers import CogVideoXPipeline + +from .cogvideox_lora import ( + calculate_noisy_latents, + collate_fn_t2v, + forward_pass, + initialize_pipeline, + load_condition_models, + load_diffusion_models, + load_latent_models, + post_latent_preparation, + prepare_conditions, + prepare_latents, + validation, +) + + +# TODO(aryan): refactor into model specs for better re-use +COGVIDEOX_T2V_FULL_FINETUNE_CONFIG = { + "pipeline_cls": CogVideoXPipeline, + "load_condition_models": load_condition_models, + "load_latent_models": load_latent_models, + "load_diffusion_models": load_diffusion_models, + "initialize_pipeline": initialize_pipeline, + "prepare_conditions": prepare_conditions, + "prepare_latents": prepare_latents, + "post_latent_preparation": post_latent_preparation, + "collate_fn": collate_fn_t2v, + "calculate_noisy_latents": calculate_noisy_latents, + "forward_pass": forward_pass, + "validation": validation, +} diff --git a/finetrainers/hunyuan_video/__init__.py b/finetrainers/hunyuan_video/__init__.py index 1800b2e..e1fdafa 100644 --- a/finetrainers/hunyuan_video/__init__.py +++ b/finetrainers/hunyuan_video/__init__.py @@ -1 +1,2 @@ -from .hunyuan_video_lora import HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG, HUNYUAN_VIDEO_T2V_LORA_CONFIG +from .full_finetune import HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG +from .hunyuan_video_lora import HUNYUAN_VIDEO_T2V_LORA_CONFIG diff --git a/finetrainers/hunyuan_video/full_finetune.py b/finetrainers/hunyuan_video/full_finetune.py new file mode 100644 index 0000000..36dd5cb --- /dev/null +++ b/finetrainers/hunyuan_video/full_finetune.py @@ -0,0 +1,30 @@ +from diffusers import HunyuanVideoPipeline + +from .hunyuan_video_lora import ( + collate_fn_t2v, + forward_pass, + initialize_pipeline, + load_condition_models, + load_diffusion_models, + load_latent_models, + post_latent_preparation, + prepare_conditions, + prepare_latents, + validation, +) + + +# TODO(aryan): refactor into model specs for better re-use +HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG = { + "pipeline_cls": HunyuanVideoPipeline, + "load_condition_models": load_condition_models, + "load_latent_models": load_latent_models, + "load_diffusion_models": load_diffusion_models, + "initialize_pipeline": initialize_pipeline, + "prepare_conditions": prepare_conditions, + "prepare_latents": prepare_latents, + "post_latent_preparation": post_latent_preparation, + "collate_fn": collate_fn_t2v, + "forward_pass": forward_pass, + "validation": validation, +} diff --git a/finetrainers/hunyuan_video/hunyuan_video_lora.py b/finetrainers/hunyuan_video/hunyuan_video_lora.py index 7d4fb5e..071451c 100644 --- a/finetrainers/hunyuan_video/hunyuan_video_lora.py +++ b/finetrainers/hunyuan_video/hunyuan_video_lora.py @@ -346,6 +346,7 @@ def _get_clip_prompt_embeds( return {"pooled_prompt_embeds": prompt_embeds} +# TODO(aryan): refactor into model specs for better re-use HUNYUAN_VIDEO_T2V_LORA_CONFIG = { "pipeline_cls": HunyuanVideoPipeline, "load_condition_models": load_condition_models, @@ -359,17 +360,3 @@ HUNYUAN_VIDEO_T2V_LORA_CONFIG = { "forward_pass": forward_pass, "validation": validation, } - -HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG = { - "pipeline_cls": HunyuanVideoPipeline, - "load_condition_models": load_condition_models, - "load_latent_models": load_latent_models, - "load_diffusion_models": load_diffusion_models, - "initialize_pipeline": initialize_pipeline, - "prepare_conditions": prepare_conditions, - "prepare_latents": prepare_latents, - "post_latent_preparation": post_latent_preparation, - "collate_fn": collate_fn_t2v, - "forward_pass": forward_pass, - "validation": validation, -} diff --git a/finetrainers/ltx_video/__init__.py b/finetrainers/ltx_video/__init__.py index 0b33512..6d5d0f9 100644 --- a/finetrainers/ltx_video/__init__.py +++ b/finetrainers/ltx_video/__init__.py @@ -1 +1,2 @@ -from .ltx_video_lora import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG, LTX_VIDEO_T2V_LORA_CONFIG +from .full_finetune import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG +from .ltx_video_lora import LTX_VIDEO_T2V_LORA_CONFIG diff --git a/finetrainers/ltx_video/full_finetune.py b/finetrainers/ltx_video/full_finetune.py new file mode 100644 index 0000000..9aa30ef --- /dev/null +++ b/finetrainers/ltx_video/full_finetune.py @@ -0,0 +1,30 @@ +from diffusers import LTXPipeline + +from .ltx_video_lora import ( + collate_fn_t2v, + forward_pass, + initialize_pipeline, + load_condition_models, + load_diffusion_models, + load_latent_models, + post_latent_preparation, + prepare_conditions, + prepare_latents, + validation, +) + + +# TODO(aryan): refactor into model specs for better re-use +LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG = { + "pipeline_cls": LTXPipeline, + "load_condition_models": load_condition_models, + "load_latent_models": load_latent_models, + "load_diffusion_models": load_diffusion_models, + "initialize_pipeline": initialize_pipeline, + "prepare_conditions": prepare_conditions, + "prepare_latents": prepare_latents, + "post_latent_preparation": post_latent_preparation, + "collate_fn": collate_fn_t2v, + "forward_pass": forward_pass, + "validation": validation, +} diff --git a/finetrainers/ltx_video/ltx_video_lora.py b/finetrainers/ltx_video/ltx_video_lora.py index 64bb625..a6e8222 100644 --- a/finetrainers/ltx_video/ltx_video_lora.py +++ b/finetrainers/ltx_video/ltx_video_lora.py @@ -322,17 +322,3 @@ LTX_VIDEO_T2V_LORA_CONFIG = { "forward_pass": forward_pass, "validation": validation, } - -LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG = { - "pipeline_cls": LTXPipeline, - "load_condition_models": load_condition_models, - "load_latent_models": load_latent_models, - "load_diffusion_models": load_diffusion_models, - "initialize_pipeline": initialize_pipeline, - "prepare_conditions": prepare_conditions, - "prepare_latents": prepare_latents, - "post_latent_preparation": post_latent_preparation, - "collate_fn": collate_fn_t2v, - "forward_pass": forward_pass, - "validation": validation, -}