From 2720d24f62cd2d4252fbe6d8ca767f36c6d8f27f Mon Sep 17 00:00:00 2001 From: CrossProduct Date: Sat, 18 Jan 2025 04:18:25 +0000 Subject: [PATCH] Moved the norm back to prepare latents step. Now its time to add the adapter layer. --- .../conditioning/condition_latents_prepare.py | 5 +---- finetrainers/ltx_video/full_finetune.py | 3 ++- finetrainers/trainer.py | 12 +----------- 3 files changed, 4 insertions(+), 16 deletions(-) diff --git a/finetrainers/conditioning/condition_latents_prepare.py b/finetrainers/conditioning/condition_latents_prepare.py index 7d4362e..ede5f36 100644 --- a/finetrainers/conditioning/condition_latents_prepare.py +++ b/finetrainers/conditioning/condition_latents_prepare.py @@ -40,8 +40,6 @@ def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int def post_conditioned_latent_patchify( latents: torch.Tensor, - latents_mean: torch.Tensor, - latents_std: torch.Tensor, num_frames: int, height: int, width: int, @@ -49,7 +47,6 @@ def post_conditioned_latent_patchify( patch_size_t: int = 1, **kwargs, ) -> torch.Tensor: - latents = _normalize_latents(latents, latents_mean, latents_std) latents = _pack_latents(latents, patch_size, patch_size_t) return {"latents": latents, "num_frames": num_frames, "height": height, "width": width} @@ -76,7 +73,7 @@ def prepare_latents_for_conditioning( latents = latents.to(dtype=dtype) _, _, num_frames, height, width = latents.shape - # latents = _normalize_latents(latents, vae.latents_mean, vae.latents_std) + latents = _normalize_latents(latents, vae.latents_mean, vae.latents_std) # latents = _pack_latents(latents, patch_size, patch_size_t) return {"latents": latents, "num_frames": num_frames, diff --git a/finetrainers/ltx_video/full_finetune.py b/finetrainers/ltx_video/full_finetune.py index ca799ea..6773070 100644 --- a/finetrainers/ltx_video/full_finetune.py +++ b/finetrainers/ltx_video/full_finetune.py @@ -11,6 +11,7 @@ from .lora import ( prepare_conditions, prepare_latents, validation, + collate_fn_t2v_cond ) @@ -24,7 +25,7 @@ LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG = { "prepare_conditions": prepare_conditions, "prepare_latents": prepare_latents, "post_latent_preparation": post_latent_preparation, - "collate_fn": collate_fn_t2v, + "collate_fn": collate_fn_t2v_cond, "forward_pass": forward_pass, "validation": validation, } diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index f38f8a6..1c3f8d3 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -799,14 +799,6 @@ class Trainer: if self.pose_condition: # condition with the pose information - # grab noise pre patchification shape - noise = torch.randn( - latent_conditions["latents"].shape, - generator=self.state.generator, - device=accelerator.device, - dtype=weight_dtype, - ) - # normalize and patchify # create noise for latent to feed diffusion transformer to get pred noise noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise @@ -819,15 +811,13 @@ class Trainer: # project this turn into patches latent_final = post_conditioned_latent_patchify(latents=latent_final, - latents_mean=(latent_conditions["mean"] + pose_video_latents["mean"] + img_refs_latents["mean"]) / 3.0, - latents_std=(latent_conditions["std"] + pose_video_latents["std"] + img_refs_latents["std"]) / 3.0, num_frames=latent_conditions["num_frames"], height=latent_conditions["height"], width=latent_conditions["width"], patch_size = 1, patch_size_t = 1) - latent_conditions.update({"noisy_latents": latent_final}) + latent_conditions.update({"noisy_latents": latent_final["latents"]}) else: # Default to flow-matching noise addition noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise