replaced the pipeline functions.

time to test this then test the weight changes.
This commit is contained in:
CrossProduct
2025-01-21 02:22:35 +00:00
parent 48d2b61029
commit 34996c1eb1
3 changed files with 26 additions and 15 deletions
@@ -36,7 +36,6 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel):
adapter_dropout=0.1, adapter_dropout=0.1,
adapter_init_scale=1e-3 adapter_init_scale=1e-3
) )
self.adapter.requires_grad_(True)
super().__init__( super().__init__(
in_channels=in_channels, in_channels=in_channels,
@@ -59,24 +58,21 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel):
@classmethod @classmethod
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
model = super().from_pretrained(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 = cls(**model.config.__dict__)
new_model.load_state_dict(model.state_dict()) new_model.load_state_dict(model.state_dict())
return new_model return new_model
def save_pretrained(self, save_directory: str, **kwargs): # def save_pretrained(self, save_directory: str, **kwargs):
""" # """
Saves model weights (including adapter) + config to disk # Saves model weights (including adapter) + config to disk
in a format compatible with .from_pretrained() # in a format compatible with .from_pretrained()
""" # """
# 1) Save config # # 1) Save config
self.config.save_pretrained(save_directory) # self.config.save_pretrained(save_directory)
# 2) Save PyTorch state dict # # 2) Save PyTorch state dict
state_dict = self.state_dict() # state_dict = self.state_dict()
torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin")) # torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin"))
def forward( def forward(
self, self,
+2 -1
View File
@@ -6,6 +6,7 @@ from .lora import (
initialize_pipeline, initialize_pipeline,
load_condition_models, load_condition_models,
load_diffusion_models, load_diffusion_models,
load_conditioned_diffusion_models,
load_latent_models, load_latent_models,
post_latent_preparation, post_latent_preparation,
prepare_conditions, prepare_conditions,
@@ -20,7 +21,7 @@ LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG = {
"pipeline_cls": LTXPipeline, "pipeline_cls": LTXPipeline,
"load_condition_models": load_condition_models, "load_condition_models": load_condition_models,
"load_latent_models": load_latent_models, "load_latent_models": load_latent_models,
"load_diffusion_models": load_diffusion_models, "load_diffusion_models": load_conditioned_diffusion_models,
"initialize_pipeline": initialize_pipeline, "initialize_pipeline": initialize_pipeline,
"prepare_conditions": prepare_conditions, "prepare_conditions": prepare_conditions,
"prepare_latents": prepare_latents, "prepare_latents": prepare_latents,
+14
View File
@@ -4,6 +4,7 @@ import torch
import torch.nn as nn import torch.nn as nn
from accelerate.logging import get_logger from accelerate.logging import get_logger
from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel
from finetrainers.conditioning import LTXVideoConditionedTransformer3DModel
from PIL import Image from PIL import Image
from transformers import T5EncoderModel, T5Tokenizer from transformers import T5EncoderModel, T5Tokenizer
@@ -37,6 +38,19 @@ def load_latent_models(
) )
return {"vae": vae} 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( def load_diffusion_models(
model_id: str = "Lightricks/LTX-Video", model_id: str = "Lightricks/LTX-Video",