diff --git a/finetrainers/conditioning/conditioned_pipeline.py b/finetrainers/conditioning/conditioned_pipeline.py index 5fae388..c13dd39 100644 --- a/finetrainers/conditioning/conditioned_pipeline.py +++ b/finetrainers/conditioning/conditioned_pipeline.py @@ -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(): diff --git a/finetrainers/ltx_video/__init__.py b/finetrainers/ltx_video/__init__.py index 69391cd..fd66532 100644 --- a/finetrainers/ltx_video/__init__.py +++ b/finetrainers/ltx_video/__init__.py @@ -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 diff --git a/finetrainers/ltx_video/full_finetune.py b/finetrainers/ltx_video/full_finetune.py index aaefdb7..241ada3 100644 --- a/finetrainers/ltx_video/full_finetune.py +++ b/finetrainers/ltx_video/full_finetune.py @@ -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, } diff --git a/finetrainers/ltx_video/full_finetune_condition.py b/finetrainers/ltx_video/full_finetune_condition.py new file mode 100644 index 0000000..f52783c --- /dev/null +++ b/finetrainers/ltx_video/full_finetune_condition.py @@ -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, +} diff --git a/finetrainers/ltx_video/lora.py b/finetrainers/ltx_video/lora.py index d9ae717..e267a40 100644 --- a/finetrainers/ltx_video/lora.py +++ b/finetrainers/ltx_video/lora.py @@ -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, diff --git a/finetrainers/models.py b/finetrainers/models.py index c24ab95..0a6da3c 100644 --- a/finetrainers/models.py +++ b/finetrainers/models.py @@ -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,