mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-18 07:33:20 +00:00
replaced the pipeline functions.
time to test this then test the weight changes.
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
Reference in New Issue
Block a user