diff --git a/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py b/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py index 72ce249..ed4c32d 100644 --- a/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py +++ b/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 - def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + @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,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, diff --git a/finetrainers/conditioning/inference_example.py b/finetrainers/conditioning/inference_example.py index 6c4ae50..d5d1eb4 100644 --- a/finetrainers/conditioning/inference_example.py +++ b/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, diff --git a/finetrainers/dataset.py b/finetrainers/dataset.py index f9d7e78..c1e8756 100644 --- a/finetrainers/dataset.py +++ b/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], diff --git a/finetrainers/ltx_video/full_finetune_condition.py b/finetrainers/ltx_video/full_finetune_condition.py index d6f3f18..a952240 100644 --- a/finetrainers/ltx_video/full_finetune_condition.py +++ b/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} diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index e9bca5a..e89565b 100644 --- a/finetrainers/trainer.py +++ b/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] + \ No newline at end of file