From 34996c1eb1bbc10d02302a609357dfcdcb72c22b Mon Sep 17 00:00:00 2001 From: CrossProduct Date: Tue, 21 Jan 2025 02:22:35 +0000 Subject: [PATCH] replaced the pipeline functions. time to test this then test the weight changes. --- .../LTXVideoConditionedTransformer3DModel.py | 24 ++++++++----------- finetrainers/ltx_video/full_finetune.py | 3 ++- finetrainers/ltx_video/lora.py | 14 +++++++++++ 3 files changed, 26 insertions(+), 15 deletions(-) diff --git a/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py b/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py index f202e07..62618ff 100644 --- a/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py +++ b/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py @@ -36,7 +36,6 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel): adapter_dropout=0.1, adapter_init_scale=1e-3 ) - self.adapter.requires_grad_(True) super().__init__( in_channels=in_channels, @@ -59,24 +58,21 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel): @classmethod def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): model = super().from_pretrained(pretrained_model_name_or_path, **kwargs) - # model is now a LTXVideoTransformer3DModel instance; - # but we want an instance of LTXVideoTransformer3DModelWithAdapter - # so we often do something like: new_model = cls(**model.config.__dict__) new_model.load_state_dict(model.state_dict()) return new_model - def save_pretrained(self, save_directory: str, **kwargs): - """ - Saves model weights (including adapter) + config to disk - in a format compatible with .from_pretrained() - """ - # 1) Save config - self.config.save_pretrained(save_directory) + # def save_pretrained(self, save_directory: str, **kwargs): + # """ + # Saves model weights (including adapter) + config to disk + # in a format compatible with .from_pretrained() + # """ + # # 1) Save config + # self.config.save_pretrained(save_directory) - # 2) Save PyTorch state dict - state_dict = self.state_dict() - torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin")) + # # 2) Save PyTorch state dict + # state_dict = self.state_dict() + # torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin")) def forward( self, diff --git a/finetrainers/ltx_video/full_finetune.py b/finetrainers/ltx_video/full_finetune.py index 6773070..ec83802 100644 --- a/finetrainers/ltx_video/full_finetune.py +++ b/finetrainers/ltx_video/full_finetune.py @@ -6,6 +6,7 @@ from .lora import ( initialize_pipeline, load_condition_models, load_diffusion_models, + load_conditioned_diffusion_models, load_latent_models, post_latent_preparation, prepare_conditions, @@ -20,7 +21,7 @@ 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, + "load_diffusion_models": load_conditioned_diffusion_models, "initialize_pipeline": initialize_pipeline, "prepare_conditions": prepare_conditions, "prepare_latents": prepare_latents, diff --git a/finetrainers/ltx_video/lora.py b/finetrainers/ltx_video/lora.py index 893d697..d244618 100644 --- a/finetrainers/ltx_video/lora.py +++ b/finetrainers/ltx_video/lora.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn from accelerate.logging import get_logger from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel +from finetrainers.conditioning import LTXVideoConditionedTransformer3DModel from PIL import Image from transformers import T5EncoderModel, T5Tokenizer @@ -37,6 +38,19 @@ def load_latent_models( ) return {"vae": vae} +def load_conditioned_diffusion_models( + model_id: str = "Lightricks/LTX-Video", + transformer_dtype: torch.dtype = torch.bfloat16, + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + **kwargs, +) -> Dict[str, nn.Module]: + transformer = LTXVideoConditionedTransformer3DModel.from_pretrained( + model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir + ) + scheduler = FlowMatchEulerDiscreteScheduler() + return {"transformer": transformer, "scheduler": scheduler} + def load_diffusion_models( model_id: str = "Lightricks/LTX-Video",