Full Finetuning for LTX possibily extended to other models.

finetrainers/args.py

@@ -455,7 +455,7 @@ def _add_training_arguments(parser: argparse.ArgumentParser) -> None:
         "--training_type",
         type=str,
         default=None,
-        help="Type of training to perform. Choose between ['lora']",
+        help="Type of training to perform. Choose between ['lora','finetune']",
     )
     parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
     parser.add_argument(

finetrainers/ltx_video/ltx_video_lora.py

@@ -321,3 +321,18 @@ def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int
     "forward_pass": forward_pass,
     "validation": validation,
 }
+
+LTX_VIDEO_T2V_FT_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,
+}
+

finetrainers/models.py

@@ -2,7 +2,7 @@
 
 from .cogvideox import COGVIDEOX_T2V_LORA_CONFIG
 from .hunyuan_video import HUNYUAN_VIDEO_T2V_LORA_CONFIG
-from .ltx_video import LTX_VIDEO_T2V_LORA_CONFIG
+from .ltx_video import LTX_VIDEO_T2V_LORA_CONFIG, LTX_VIDEO_T2V_FT_CONFIG
 
 
 SUPPORTED_MODEL_CONFIGS = {
@@ -11,6 +11,7 @@
     },
     "ltx_video": {
         "lora": LTX_VIDEO_T2V_LORA_CONFIG,
+        "finetune": LTX_VIDEO_T2V_FT_CONFIG,
     },
     "cogvideox": {
         "lora": COGVIDEOX_T2V_LORA_CONFIG,

finetrainers/trainer.py

@@ -99,6 +99,8 @@ def __init__(self, args: Args) -> None:
         self.state.model_name = self.args.model_name
         self.model_config = get_config_from_model_name(self.args.model_name, self.args.training_type)
 
+        # Components list
+        self.components = []
     def prepare_dataset(self) -> None:
         # TODO(aryan): Make a background process for fetching
         logger.info("Initializing dataset and dataloader")
@@ -153,6 +155,17 @@ def _set_components(self, components: Dict[str, Any]) -> None:
         self.transformer_config = self.transformer.config if self.transformer is not None else self.transformer_config
         self.vae_config = self.vae.config if self.vae is not None else self.vae_config
 
+        self.components = [self.tokenizer, 
+                           self.tokenizer_2, 
+                           self.tokenizer_3, 
+                           self.text_encoder, 
+                           self.text_encoder_2, 
+                           self.text_encoder_3, 
+                           self.transformer, 
+                           self.unet, 
+                           self.vae]
+
+
     def _delete_components(self) -> None:
         self.tokenizer = None
         self.tokenizer_2 = None
@@ -167,6 +180,8 @@ def _delete_components(self) -> None:
         free_memory()
         torch.cuda.synchronize(self.state.accelerator.device)
 
+        self.components = None
+
     def prepare_models(self) -> None:
         logger.info("Initializing models")
 
@@ -189,6 +204,16 @@ def prepare_models(self) -> None:
             if self.args.enable_tiling:
                 self.vae.enable_tiling()
 
+    def _disable_grad_for_components(self, components:list):
+        for component in components:
+            if component is not None:
+                component.requires_grad_(False)
+
+    def _enable_grad_for_components(self, components:list):
+        for component in components:
+            if component is not None:
+                component.requires_grad_(True)
+
     def prepare_precomputations(self) -> None:
         if not self.args.precompute_conditions:
             return
@@ -237,16 +262,11 @@ def collate_fn(batch):
         self._set_components(condition_components)
         self._move_components_to_device()
 
-        # TODO(aryan): refactor later. for now only lora is supported
-        components_to_disable_grads = [
+        self._disable_grad_for_components(components=[
             self.text_encoder,
             self.text_encoder_2,
             self.text_encoder_3,
-        ]
-        for component in components_to_disable_grads:
-            if component is not None:
-                component.requires_grad_(False)
-
+        ])
         if self.args.caption_dropout_p > 0 and self.args.caption_dropout_technique == "empty":
             logger.warning(
                 "Caption dropout is not supported with precomputation yet. This will be supported in the future."
@@ -300,12 +320,7 @@ def collate_fn(batch):
         self._set_components(latent_components)
         self._move_components_to_device()
 
-        # TODO(aryan): refactor later
-        components_to_disable_grads = [self.vae]
-        for component in components_to_disable_grads:
-            if component is not None:
-                component.requires_grad_(False)
-
+        self._disable_grad_for_components(components=[self.vae])
         if self.vae is not None:
             if self.args.enable_slicing:
                 self.vae.enable_slicing()
@@ -363,17 +378,18 @@ def prepare_trainable_parameters(self) -> None:
         diffusion_components = self.model_config["load_diffusion_models"](**self._get_load_components_kwargs())
         self._set_components(diffusion_components)
 
-        # TODO(aryan): refactor later. for now only lora is supported
-        components_to_disable_grads = [
+        self._disable_grad_for_components(components=[
             self.text_encoder,
             self.text_encoder_2,
             self.text_encoder_3,
-            self.transformer,
             self.vae,
-        ]
-        for component in components_to_disable_grads:
-            if component is not None:
-                component.requires_grad_(False)
+        ])
+
+        if self.args.training_type == "full_finetune":
+            logger.info("Full Fine Tuning Enabled")
+            self._enable_grad_for_components(components=[self.transformer])
+        else:
+            logger.info("Lora Fine Tuning Enabled")
 
         # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision
         # as these weights are only used for inference, keeping weights in full precision is not required.
@@ -398,13 +414,14 @@ def prepare_trainable_parameters(self) -> None:
         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)
 
         # 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():