storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-13 23:53:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reducing cogload on the framework.

Browse Source
This commit is contained in:
CrossProduct
2025-01-24 22:46:32 +00:00
parent 305a4a994d
commit c6490c0d51
6 changed files with 158 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/conditioning/conditioned_pipeline.py
+4
View File
@@ -13,6 +13,10 @@ from typing import Callable, Any
import numpy as np
from diffusers.utils import is_torch_xla_available
from diffusers.pipelines.ltx.pipeline_output import LTXPipelineOutput
from diffusers.utils import is_torch_xla_available
from torch import torch
from transformers import T5EncoderModel, T5TokenizerFast
from typing import Any, Callable, Dict, List, Optional, Union
if is_torch_xla_available():
finetrainers/ltx_video/__init__.py
+1
View File
@@ -1,2 +1,3 @@
from .full_finetune import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG
from .full_finetune_condition import LTX_VIDEO_COND_T2V_FULL_FINETUNE_CONFIG
from .lora import LTX_VIDEO_T2V_LORA_CONFIG
finetrainers/ltx_video/full_finetune.py
+5 -9
View File
@@ -1,4 +1,5 @@
from diffusers import LTXPipeline
from finetrainers.conditioning.conditioned_pipeline import LTXConditionedPipeline
from .lora import (
collate_fn_t2v,
@@ -6,29 +7,24 @@ from .lora import (
initialize_pipeline,
load_condition_models,
load_diffusion_models,
load_conditioned_diffusion_models, # use this for conditioning
load_latent_models,
post_latent_preparation,
prepare_conditions,
prepare_latents,
validation,
conditional_validation,
collate_fn_t2v_cond,
conditioned_forward_pass # use this for conditioning
)
# TODO(aryan): refactor into model specs for better re-use
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_conditioned_diffusion_models,
"load_diffusion_models": load_diffusion_models,
"initialize_pipeline": initialize_pipeline,
"prepare_conditions": prepare_conditions,
"prepare_latents": prepare_latents,
"post_latent_preparation": post_latent_preparation,
"collate_fn": collate_fn_t2v_cond,
"forward_pass": conditioned_forward_pass,
"validation": conditional_validation,
"collate_fn": collate_fn_t2v,
"forward_pass": forward_pass,
"validation": validation,
}
finetrainers/ltx_video/full_finetune_condition.py
+141
View File
@@ -0,0 +1,141 @@
from typing import Dict, List, Optional
import torch
import torch.nn as nn
from diffusers import FlowMatchEulerDiscreteScheduler, LTXVideoTransformer3DModel
from finetrainers.conditioning.LTXVideoConditionedTransformer3DModel import LTXVideoConditionedTransformer3DModel
from finetrainers.conditioning.conditioned_pipeline import LTXConditionedPipeline
from PIL import Image
# Exisiting Pipeline
from .lora import (
initialize_pipeline,
load_conditioned_diffusion_models, # use this for conditioning
load_latent_models,
post_latent_preparation,
prepare_conditions,
prepare_latents,
load_condition_models
)
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",
transformer_dtype: torch.dtype = torch.bfloat16,
revision: Optional[str] = None,
cache_dir: Optional[str] = None,
**kwargs,
) -> Dict[str, nn.Module]:
transformer = LTXVideoTransformer3DModel.from_pretrained(
model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir
)
scheduler = FlowMatchEulerDiscreteScheduler()
return {"transformer": transformer, "scheduler": scheduler}
def collate_fn_t2v_cond(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torch.Tensor]:
return {
"prompts": [x["prompt"] for x in batch[0]],
"videos": torch.stack([x["video"] for x in batch[0]]),
"poses": torch.stack([x["pose"] for x in batch[0]]),
"img_refs": torch.stack([x["img_ref"] for x in batch[0]])
}
def conditioned_forward_pass(
transformer: LTXVideoConditionedTransformer3DModel,
prompt_embeds: torch.Tensor,
prompt_attention_mask: torch.Tensor,
latents: torch.Tensor,
noisy_latents: torch.Tensor,
timesteps: torch.LongTensor,
num_frames: int,
height: int,
width: int,
noisy_latents_residual:torch.Tensor,
**kwargs,
) -> torch.Tensor:
rope_interpolation_scale = [1 / 25, 32, 32]
denoised_latents = transformer(
hidden_states=noisy_latents,
encoder_hidden_states=prompt_embeds,
timestep=timesteps,
encoder_attention_mask=prompt_attention_mask,
num_frames=num_frames,
height=height,
width=width,
rope_interpolation_scale=rope_interpolation_scale,
return_dict=False,
residual_x=noisy_latents_residual,
)[0]
return {"latents": denoised_latents}
def conditional_validation(
pipeline: LTXConditionedPipeline,
prompt: str,
image: Optional[Image.Image] = None,
video: Optional[List[Image.Image]] = None,
pose_video: Optional[List[Image.Image]] = None,
image_ref_video: Optional[List[Image.Image]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_frames: Optional[int] = None,
frame_rate: int = 24,
num_videos_per_prompt: int = 1,
generator: Optional[torch.Generator] = None,
**kwargs,
):
generation_kwargs = {
"prompt": prompt,
"height": height,
"width": width,
"num_frames": num_frames,
"frame_rate": frame_rate,
"num_videos_per_prompt": num_videos_per_prompt,
"generator": generator,
"return_dict": True,
"output_type": "pil",
"pose_video": pose_video,
"image_ref_video": image_ref_video,
}
# pose template noisey input [cat] img_ref + pose video
# pose template ref video latent + patchify video latent.
# img ref patchify video latent +
# add them together as input
# residual x latent
generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None}
video = pipeline(**generation_kwargs).frames[0]
return [("video", video)]
LTX_VIDEO_COND_T2V_FULL_FINETUNE_CONFIG = {
"pipeline_cls": LTXConditionedPipeline,
"load_condition_models": load_condition_models,
"load_latent_models": load_latent_models,
"load_diffusion_models": load_conditioned_diffusion_models,
"initialize_pipeline": initialize_pipeline,
"prepare_conditions": prepare_conditions,
"prepare_latents": prepare_latents,
"post_latent_preparation": post_latent_preparation,
"collate_fn": collate_fn_t2v_cond,
"forward_pass": conditioned_forward_pass,
"validation": conditional_validation,
}
finetrainers/ltx_video/lora.py
+5 -81
View File
@@ -5,13 +5,12 @@ import torch.nn as nn
from accelerate.logging import get_logger
from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel
from finetrainers.conditioning.LTXVideoConditionedTransformer3DModel import LTXVideoConditionedTransformer3DModel
from finetrainers.conditioning.conditioned_pipeline import LTXConditionedPipeline
from PIL import Image
from transformers import T5EncoderModel, T5Tokenizer
logger = get_logger("finetrainers") # pylint: disable=invalid-name
def load_condition_models(
model_id: str = "Lightricks/LTX-Video",
text_encoder_dtype: torch.dtype = torch.bfloat16,
@@ -38,9 +37,9 @@ def load_latent_models(
)
return {"vae": vae}
def load_conditioned_diffusion_models(
def load_condition_models(
model_id: str = "Lightricks/LTX-Video",
transformer_dtype: torch.dtype = torch.bfloat16,
text_encoder_dtype: torch.dtype = torch.bfloat16,
revision: Optional[str] = None,
cache_dir: Optional[str] = None,
**kwargs,
@@ -84,7 +83,7 @@ def initialize_pipeline(
enable_tiling: bool = False,
enable_model_cpu_offload: bool = False,
**kwargs,
) -> LTXPipeline:
) -> LTXConditionedPipeline:
component_name_pairs = [
("tokenizer", tokenizer),
("text_encoder", text_encoder),
@@ -97,7 +96,7 @@ def initialize_pipeline(
if component is not None:
components[name] = component
pipe = LTXPipeline.from_pretrained(model_id, **components, revision=revision, cache_dir=cache_dir)
pipe = LTXConditionedPipeline.from_pretrained(model_id, **components, revision=revision, cache_dir=cache_dir)
pipe.text_encoder = pipe.text_encoder.to(dtype=text_encoder_dtype)
pipe.transformer = pipe.transformer.to(dtype=transformer_dtype)
pipe.vae = pipe.vae.to(dtype=vae_dtype)
@@ -196,14 +195,6 @@ def post_latent_preparation(
return {"latents": latents, "num_frames": num_frames, "height": height, "width": width}
def collate_fn_t2v_cond(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torch.Tensor]:
return {
"prompts": [x["prompt"] for x in batch[0]],
"videos": torch.stack([x["video"] for x in batch[0]]),
"poses": torch.stack([x["pose"] for x in batch[0]]),
"img_refs": torch.stack([x["img_ref"] for x in batch[0]])
}
def collate_fn_t2v(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torch.Tensor]:
return {
"prompts": [x["prompt"] for x in batch[0]],
@@ -211,35 +202,6 @@ def collate_fn_t2v(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torc
}
def conditioned_forward_pass(
transformer: LTXVideoConditionedTransformer3DModel,
prompt_embeds: torch.Tensor,
prompt_attention_mask: torch.Tensor,
latents: torch.Tensor,
noisy_latents: torch.Tensor,
timesteps: torch.LongTensor,
num_frames: int,
height: int,
width: int,
noisy_latents_residual:torch.Tensor,
**kwargs,
) -> torch.Tensor:
rope_interpolation_scale = [1 / 25, 32, 32]
denoised_latents = transformer(
hidden_states=noisy_latents,
encoder_hidden_states=prompt_embeds,
timestep=timesteps,
encoder_attention_mask=prompt_attention_mask,
num_frames=num_frames,
height=height,
width=width,
rope_interpolation_scale=rope_interpolation_scale,
return_dict=False,
residual_x=noisy_latents_residual,
)[0]
return {"latents": denoised_latents}
def forward_pass(
transformer: LTXVideoTransformer3DModel,
@@ -270,45 +232,7 @@ def forward_pass(
return {"latents": denoised_latents}
def conditional_validation(
pipeline: LTXPipeline,
prompt: str,
image: Optional[Image.Image] = None,
video: Optional[List[Image.Image]] = None,
pose_video: Optional[List[Image.Image]] = None,
image_ref_video: Optional[List[Image.Image]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_frames: Optional[int] = None,
frame_rate: int = 24,
num_videos_per_prompt: int = 1,
generator: Optional[torch.Generator] = None,
**kwargs,
):
generation_kwargs = {
"prompt": prompt,
"height": height,
"width": width,
"num_frames": num_frames,
"frame_rate": frame_rate,
"num_videos_per_prompt": num_videos_per_prompt,
"generator": generator,
"return_dict": True,
"output_type": "pil",
"pose_video": pose_video,
"image_ref_video": image_ref_video,
}
# pose template noisey input [cat] img_ref + pose video
# pose template ref video latent + patchify video latent.
# img ref patchify video latent +
# add them together as input
# residual x latent
generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None}
video = pipeline(**generation_kwargs).frames[0]
def validation(
pipeline: LTXPipeline,
prompt: str,
finetrainers/models.py
+2 -2
View File
@@ -2,7 +2,7 @@ from typing import Any, Dict
from .cogvideox import COGVIDEOX_T2V_FULL_FINETUNE_CONFIG, COGVIDEOX_T2V_LORA_CONFIG
from .hunyuan_video import HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG, HUNYUAN_VIDEO_T2V_LORA_CONFIG
from .ltx_video import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG, LTX_VIDEO_T2V_LORA_CONFIG
from .ltx_video import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG,LTX_VIDEO_COND_T2V_FULL_FINETUNE_CONFIG, LTX_VIDEO_T2V_LORA_CONFIG
SUPPORTED_MODEL_CONFIGS = {
@@ -12,7 +12,7 @@ SUPPORTED_MODEL_CONFIGS = {
},
"ltx_video": {
"lora": LTX_VIDEO_T2V_LORA_CONFIG,
"full-finetune": LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG,
"full-finetune":LTX_VIDEO_COND_T2V_FULL_FINETUNE_CONFIG # LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG,
},
"cogvideox": {
"lora": COGVIDEOX_T2V_LORA_CONFIG,