From b7f26bd2e680f250484fb0c4aa2696eb05e62a7e Mon Sep 17 00:00:00 2001 From: CrossProduct Date: Thu, 23 Jan 2025 23:49:24 +0000 Subject: [PATCH] removed some code made things clearer to replace the pipeline inference which is needed for validation. --- .../conditioning/condition_latents_prepare.py | 4 +- finetrainers/ltx_video/full_finetune.py | 7 ++-- finetrainers/ltx_video/lora.py | 36 +++++++++++++++++ finetrainers/trainer.py | 40 ++++++++++--------- 4 files changed, 63 insertions(+), 24 deletions(-) diff --git a/finetrainers/conditioning/condition_latents_prepare.py b/finetrainers/conditioning/condition_latents_prepare.py index ede5f36..9561ef4 100644 --- a/finetrainers/conditioning/condition_latents_prepare.py +++ b/finetrainers/conditioning/condition_latents_prepare.py @@ -78,6 +78,4 @@ def prepare_latents_for_conditioning( return {"latents": latents, "num_frames": num_frames, "height": height, - "width": width, - "mean":vae.latents_mean, - "std":vae.latents_std} \ No newline at end of file + "width": width } \ No newline at end of file diff --git a/finetrainers/ltx_video/full_finetune.py b/finetrainers/ltx_video/full_finetune.py index b5553c9..aaefdb7 100644 --- a/finetrainers/ltx_video/full_finetune.py +++ b/finetrainers/ltx_video/full_finetune.py @@ -6,14 +6,15 @@ from .lora import ( initialize_pipeline, load_condition_models, load_diffusion_models, - load_conditioned_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 + conditioned_forward_pass # use this for conditioning ) @@ -29,5 +30,5 @@ LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG = { "post_latent_preparation": post_latent_preparation, "collate_fn": collate_fn_t2v_cond, "forward_pass": conditioned_forward_pass, - "validation": validation, + "validation": conditional_validation, } diff --git a/finetrainers/ltx_video/lora.py b/finetrainers/ltx_video/lora.py index 8f33da0..dec3173 100644 --- a/finetrainers/ltx_video/lora.py +++ b/finetrainers/ltx_video/lora.py @@ -66,6 +66,7 @@ def load_diffusion_models( return {"transformer": transformer, "scheduler": scheduler} + def initialize_pipeline( model_id: str = "Lightricks/LTX-Video", text_encoder_dtype: torch.dtype = torch.bfloat16, @@ -269,7 +270,42 @@ 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, + 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", + "residual_x": None, + } + + # 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/trainer.py b/finetrainers/trainer.py index c4267c3..45863f4 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -676,6 +676,7 @@ class Trainer: if random.random() < self.args.caption_dropout_p: prompts = [""] * batch_size if self.pose_condition == True: + # this is vae output of target video latent_conditions = condition_latents_prepare.prepare_latents_for_conditioning( vae=self.vae, image_or_video=videos, @@ -726,7 +727,7 @@ class Trainer: latent_conditions = make_contiguous(latent_conditions) if self.pose_condition: - # VAE output not patchified yet. + # Pose Template VAE. pose_video_latents = condition_latents_prepare.prepare_latents_for_conditioning( vae=self.vae, image_or_video=poses, @@ -739,6 +740,7 @@ class Trainer: pose_video_latents = make_contiguous(pose_video_latents) + # Image Ref VAE. img_refs_latents = condition_latents_prepare.prepare_latents_for_conditioning( vae=self.vae, image_or_video=img_refs, @@ -751,8 +753,6 @@ class Trainer: img_refs_latents = make_contiguous(img_refs_latents) - - text_conditions = make_contiguous(text_conditions) if self.args.caption_dropout_technique == "zero": @@ -776,13 +776,14 @@ class Trainer: device=accelerator.device, generator=self.state.generator, ) + timesteps = (sigmas * 1000.0).long() noise = torch.randn( - latent_conditions["latents"].shape, - generator=self.state.generator, - device=accelerator.device, - dtype=weight_dtype, + latent_conditions["latents"].shape, + generator=self.state.generator, + device=accelerator.device, + dtype=weight_dtype, ) sigmas = expand_tensor_dims(sigmas, ndim=noise.ndim) @@ -797,38 +798,41 @@ class Trainer: timesteps=timesteps, ) else: - if self.pose_condition: + # Note in pose mode all shapes are VAE output shape and not patchified until the end # condition with the pose information # 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 - noisy_latents = noisy_latents.to(latent_conditions["latents"].dtype) + target_video_latents = latent_conditions["latents"] + noisy_latents = (1.0 - sigmas) * target_video_latents + sigmas * noise + noisy_latents = noisy_latents.to(target_video_latents.dtype) + # add pose information to both channels pose_noisy_latents = noisy_latents + pose_video_latents["latents"] pose_img_ref_latents = img_refs_latents["latents"] + pose_video_latents["latents"] - # expand channel information # B x 2C latent will be projected to adapter to scale it back to 128d - latent_final = torch.cat([pose_img_ref_latents,pose_noisy_latents], dim=1) + + # expand channel information # B x 2C latent will be projected to adapter to scale it back to 128d using adapter. + pose_img_ref_latents = torch.cat([pose_img_ref_latents,pose_noisy_latents], dim=1) # project this turn into patches - latent_final = post_conditioned_latent_patchify(latents=latent_final, + pose_img_ref_latents = post_conditioned_latent_patchify(latents=pose_img_ref_latents, num_frames=latent_conditions["num_frames"], height=latent_conditions["height"], width=latent_conditions["width"], patch_size = 1, patch_size_t = 1) - + # target video as a input residual noisy_latents_residual = post_conditioned_latent_patchify(latents=noisy_latents, num_frames=latent_conditions["num_frames"], height=latent_conditions["height"], width=latent_conditions["width"], patch_size = 1, patch_size_t = 1) - # target latent patchified - - # pose information latent - latent_conditions.update({"noisy_latents": latent_final["latents"]}) + # Target Latent Patchified + # pose template noisey input [cat] img_ref + pose video + latent_conditions.update({"noisy_latents": pose_img_ref_latents["latents"]}) + # input video noise at level residual information to adapter latent_conditions.update({"noisy_latents_residual":noisy_latents_residual["latents"]}) else: # Default to flow-matching noise addition