Adapter weights can be loaded.

finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py

@@ -1,6 +1,6 @@
 import math
 from typing import Any, Dict, Optional, Tuple
-
+import os
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -57,8 +57,8 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel):
             adapter_init_scale=1e-3
         )
 
-    @classmethod
+    @classmethod # rewrite this to ensure that if it doesn't have adapter weights it will load this way but if it has adapter weights it will not.
-    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+    def from_pretrained(cls, pretrained_model_name_or_path,save_directory, **kwargs):
         # First create an empty model with the desired architecture
         print("Pretrain Loading")
         model = cls()
@@ -73,12 +73,34 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel):
         filtered_dict = {}
         for k, v in pretrained_dict.items():
             if k in model_dict:
+                
                 filtered_dict[k] = v
+            else:
+                print(k)
         # Update model with pretrained weights
         model_dict.update(filtered_dict)
+
+        adapter_weights_path = os.path.join(save_directory, "adapter_weights.pth")
+        if os.path.exists(adapter_weights_path):
+            adapter_weights = torch.load(adapter_weights_path)
+            model.adapter.load_state_dict(adapter_weights)
+            print("Adapter weights loaded successfully.")
+        else:
+            print("No adapter weights file found. Using default adapter weights.")
+
         model.load_state_dict(model_dict)
         return model
 
+    def save_pretrained(self,save_directory):
+        model_state_dict = self.state_dict()
+        torch.save(model_state_dict, save_directory)
+        
+        adapter_state_dict = self.adapter.state_dict()
+        save_directory = os.path.dirname(save_directory)
+        path = os.path.join(save_directory, "adapter_weights.pth")
+        torch.save(adapter_state_dict, path)
+        print("Saving Model with Adapter")
+
     def forward(
             self,
             hidden_states: torch.Tensor,

finetrainers/conditioning/inference_example.py

@@ -5,7 +5,7 @@ from diffusers.utils import export_to_video
 from diffusers.utils import testing_utils
 testing_utils.enable_full_determinism()
 
-transformer = LTXVideoTransformer3DModel.from_pretrained("/home/ubuntu/upstream/finetrainers/ltx-video/ltxv_strip/checkpoint-38000/transformer/",torch_dtype=torch.bfloat16).to("cuda")
+#transformer = LTXVideoTransformer3DModel.from_pretrained("/home/ubuntu/upstream/finetrainers/ltx-video/ltxv_strip/checkpoint-38000/transformer/",torch_dtype=torch.bfloat16).to("cuda")
 pipe = LTXPipeline.from_pretrained(
     "Lightricks/LTX-Video", torch_dtype=torch.bfloat16,
     # transformer=transformer
@@ -20,6 +20,7 @@ for seconds in range(3, 8):
     num_frames = seconds * 24
 
     result = pipe(
+        guidance_scale=0,
         prompt=prompt,
         width=768,
         height=512,

finetrainers/dataset.py

@@ -299,6 +299,8 @@ class ImageOrVideoDataset(Dataset):
         nearest_res = self._find_nearest_resolution(frames.shape[2], frames.shape[3])
         frames_resized = torch.stack([resize(frame, nearest_res) for frame in frames], dim=0)
         frames = torch.stack([self.video_transforms(frame) for frame in frames_resized], dim=0)
+        # torch.Size([96, 3, 512, 768])
+    
 
         return frames
 
@@ -386,6 +388,7 @@ class ImageOrVideoDatasetWithResizeAndRectangleCrop(ImageOrVideoDataset):
 
     def _preprocess_video(self, path: Path) -> torch.Tensor:
         video_reader = decord.VideoReader(uri=path.as_posix())
+        
         video_num_frames = len(video_reader)
         nearest_frame_bucket = min(
             [bucket for bucket in self.resolution_buckets if bucket <= video_num_frames],

finetrainers/ltx_video/full_finetune_condition.py

@@ -73,7 +73,8 @@ def load_conditioned_diffusion_models(
     **kwargs,
 ) -> Dict[str, nn.Module]:
     transformer = LTXVideoConditionedTransformer3DModel.from_pretrained(
-        model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir
+        model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir,
+        save_directory=kwargs["output_dir"]
     )
     scheduler = FlowMatchEulerDiscreteScheduler()
     return {"transformer": transformer, "scheduler": scheduler}

finetrainers/trainer.py

@@ -330,7 +330,7 @@ class Trainer:
     def prepare_trainable_parameters(self) -> None:
         logger.info("Initializing trainable parameters")
 
-        diffusion_components = self.model_config["load_diffusion_models"](**self._get_load_components_kwargs())
+        diffusion_components = self.model_config["load_diffusion_models"](output_dir=self.args.output_dir,**self._get_load_components_kwargs())
         self._set_components(diffusion_components)
 
         components = [self.text_encoder, self.text_encoder_2, self.text_encoder_3, self.vae]
@@ -1039,9 +1039,9 @@ class Trainer:
             # loading videos for inference ..
             if self.pose_condition:
                 if pose_video is not None:
-                    pose_video = self.preprocess_condition_video(pose_video,max_num_frames=self.dataset.max_num_frames)
+                    pose_video = self.dataset._preprocess_video(Path(pose_video)).unsqueeze(0) #self.preprocess_condition_video(pose_video,max_num_frames=self.dataset.max_num_frames)
                 if img_ref_video is not None:
-                    img_ref_video = self.preprocess_condition_video_image_reference_video(path=img_ref_video,resolution_buckets=self.dataset.resolution_buckets,max_num_frames=self.dataset.max_num_frames)
+                    img_ref_video = self.dataset._preprocess_video_image_reference_video(Path(img_ref_video)).unsqueeze(0) #self.preprocess_condition_video_image_reference_video(path=img_ref_video,resolution_buckets=self.dataset.resolution_buckets,max_num_frames=self.dataset.max_num_frames)
                 
             logger.debug(
                 f"Validating sample {i + 1}/{num_validation_samples} on process {accelerator.process_index}. Prompt: {prompt}",
@@ -1393,3 +1393,9 @@ class Trainer:
         # adding singleton dimension for downstream inference.
         frames = frames.unsqueeze(0)
         return frames
+    
+    @staticmethod
+    def _find_nearest_resolution(resolutions, height, width):
+        nearest_res = min(resolutions, key=lambda x: abs(x[1] - height) + abs(x[2] - width))
+        return nearest_res[1], nearest_res[2]
+