upgrade finetrainers

finetrainers/args.py

@@ -447,7 +447,7 @@ class BaseArgs:
         }
 
         training_arguments = {
-            "training_type":self.training_type,
+            "training_type": self.training_type,
             "seed": self.seed,
             "batch_size": self.batch_size,
             "train_steps": self.train_steps,

finetrainers/patches/__init__.py

@@ -17,7 +17,12 @@ def perform_patches_for_training(args: "BaseArgs", parallel_backend: "ParallelBa
         if parallel_backend.tensor_parallel_enabled:
             patch.patch_apply_rotary_emb_for_tp_compatibility()
 
+    if args.model_name == ModelType.WAN:
+        from .models.wan import patch
+
+        patch.patch_time_text_image_embedding_forward()
+
     if args.training_type == TrainingType.LORA and len(args.layerwise_upcasting_modules) > 0:
-        from dependencies.peft import patch
+        from .dependencies.peft import patch
 
         patch.patch_peft_move_adapter_to_device_of_base_layer()

finetrainers/patches/models/ltx_video/patch.py

@@ -16,7 +16,7 @@ def patch_apply_rotary_emb_for_tp_compatibility() -> None:
 
 
 def _perform_ltx_transformer_forward_patch() -> None:
-    LTXVideoTransformer3DModel.forward = _patched_LTXVideoTransformer3Dforward
+    LTXVideoTransformer3DModel.forward = _patched_LTXVideoTransformer3D_forward
 
 
 def _perform_ltx_apply_rotary_emb_tensor_parallel_compatibility_patch() -> None:
@@ -35,7 +35,7 @@ def _perform_ltx_apply_rotary_emb_tensor_parallel_compatibility_patch() -> None:
     diffusers.models.transformers.transformer_ltx.apply_rotary_emb = apply_rotary_emb
 
 
-def _patched_LTXVideoTransformer3Dforward(
+def _patched_LTXVideoTransformer3D_forward(
     self,
     hidden_states: torch.Tensor,
     encoder_hidden_states: torch.Tensor,

finetrainers/patches/models/wan/patch.py

@@ -0,0 +1,33 @@
+from typing import Optional
+
+import diffusers
+import torch
+
+
+def patch_time_text_image_embedding_forward() -> None:
+    _patch_time_text_image_embedding_forward()
+
+
+def _patch_time_text_image_embedding_forward() -> None:
+    diffusers.models.transformers.transformer_wan.WanTimeTextImageEmbedding.forward = (
+        _patched_WanTimeTextImageEmbedding_forward
+    )
+
+
+def _patched_WanTimeTextImageEmbedding_forward(
+    self,
+    timestep: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    encoder_hidden_states_image: Optional[torch.Tensor] = None,
+):
+    # Some code has been removed compared to original implementation in Diffusers
+    # Also, timestep is typed as that of encoder_hidden_states
+    timestep = self.timesteps_proj(timestep).type_as(encoder_hidden_states)
+    temb = self.time_embedder(timestep).type_as(encoder_hidden_states)
+    timestep_proj = self.time_proj(self.act_fn(temb))
+
+    encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+    if encoder_hidden_states_image is not None:
+        encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image)
+
+    return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image

finetrainers/trainer/sft_trainer/trainer.py

@@ -334,6 +334,7 @@ class SFTTrainer:
         parallel_backend = self.state.parallel_backend
         train_state = self.state.train_state
         device = parallel_backend.device
+        dtype = self.args.transformer_dtype
 
         memory_statistics = utils.get_memory_statistics()
         logger.info(f"Memory before training start: {json.dumps(memory_statistics, indent=4)}")
@@ -447,8 +448,8 @@ class SFTTrainer:
 
             logger.debug(f"Starting training step ({train_state.step}/{self.args.train_steps})")
 
-            utils.align_device_and_dtype(latent_model_conditions, device, self.args.transformer_dtype)
-            utils.align_device_and_dtype(condition_model_conditions, device, self.args.transformer_dtype)
+            latent_model_conditions = utils.align_device_and_dtype(latent_model_conditions, device, dtype)
+            condition_model_conditions = utils.align_device_and_dtype(condition_model_conditions, device, dtype)
             latent_model_conditions = utils.make_contiguous(latent_model_conditions)
             condition_model_conditions = utils.make_contiguous(condition_model_conditions)