removed some code made things clearer to replace the pipeline inference which is needed for validation.

This commit is contained in:
CrossProduct
2025-01-23 23:49:24 +00:00
parent d6b3615421
commit b7f26bd2e6
4 changed files with 63 additions and 24 deletions
@@ -78,6 +78,4 @@ def prepare_latents_for_conditioning(
return {"latents": latents, return {"latents": latents,
"num_frames": num_frames, "num_frames": num_frames,
"height": height, "height": height,
"width": width, "width": width }
"mean":vae.latents_mean,
"std":vae.latents_std}
+4 -3
View File
@@ -6,14 +6,15 @@ from .lora import (
initialize_pipeline, initialize_pipeline,
load_condition_models, load_condition_models,
load_diffusion_models, load_diffusion_models,
load_conditioned_diffusion_models, load_conditioned_diffusion_models, # use this for conditioning
load_latent_models, load_latent_models,
post_latent_preparation, post_latent_preparation,
prepare_conditions, prepare_conditions,
prepare_latents, prepare_latents,
validation, validation,
conditional_validation,
collate_fn_t2v_cond, 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, "post_latent_preparation": post_latent_preparation,
"collate_fn": collate_fn_t2v_cond, "collate_fn": collate_fn_t2v_cond,
"forward_pass": conditioned_forward_pass, "forward_pass": conditioned_forward_pass,
"validation": validation, "validation": conditional_validation,
} }
+36
View File
@@ -66,6 +66,7 @@ def load_diffusion_models(
return {"transformer": transformer, "scheduler": scheduler} return {"transformer": transformer, "scheduler": scheduler}
def initialize_pipeline( def initialize_pipeline(
model_id: str = "Lightricks/LTX-Video", model_id: str = "Lightricks/LTX-Video",
text_encoder_dtype: torch.dtype = torch.bfloat16, text_encoder_dtype: torch.dtype = torch.bfloat16,
@@ -269,7 +270,42 @@ def forward_pass(
return {"latents": denoised_latents} 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( def validation(
pipeline: LTXPipeline, pipeline: LTXPipeline,
prompt: str, prompt: str,
+22 -18
View File
@@ -676,6 +676,7 @@ class Trainer:
if random.random() < self.args.caption_dropout_p: if random.random() < self.args.caption_dropout_p:
prompts = [""] * batch_size prompts = [""] * batch_size
if self.pose_condition == True: if self.pose_condition == True:
# this is vae output of target video
latent_conditions = condition_latents_prepare.prepare_latents_for_conditioning( latent_conditions = condition_latents_prepare.prepare_latents_for_conditioning(
vae=self.vae, vae=self.vae,
image_or_video=videos, image_or_video=videos,
@@ -726,7 +727,7 @@ class Trainer:
latent_conditions = make_contiguous(latent_conditions) latent_conditions = make_contiguous(latent_conditions)
if self.pose_condition: if self.pose_condition:
# VAE output not patchified yet. # Pose Template VAE.
pose_video_latents = condition_latents_prepare.prepare_latents_for_conditioning( pose_video_latents = condition_latents_prepare.prepare_latents_for_conditioning(
vae=self.vae, vae=self.vae,
image_or_video=poses, image_or_video=poses,
@@ -739,6 +740,7 @@ class Trainer:
pose_video_latents = make_contiguous(pose_video_latents) pose_video_latents = make_contiguous(pose_video_latents)
# Image Ref VAE.
img_refs_latents = condition_latents_prepare.prepare_latents_for_conditioning( img_refs_latents = condition_latents_prepare.prepare_latents_for_conditioning(
vae=self.vae, vae=self.vae,
image_or_video=img_refs, image_or_video=img_refs,
@@ -751,8 +753,6 @@ class Trainer:
img_refs_latents = make_contiguous(img_refs_latents) img_refs_latents = make_contiguous(img_refs_latents)
text_conditions = make_contiguous(text_conditions) text_conditions = make_contiguous(text_conditions)
if self.args.caption_dropout_technique == "zero": if self.args.caption_dropout_technique == "zero":
@@ -776,13 +776,14 @@ class Trainer:
device=accelerator.device, device=accelerator.device,
generator=self.state.generator, generator=self.state.generator,
) )
timesteps = (sigmas * 1000.0).long() timesteps = (sigmas * 1000.0).long()
noise = torch.randn( noise = torch.randn(
latent_conditions["latents"].shape, latent_conditions["latents"].shape,
generator=self.state.generator, generator=self.state.generator,
device=accelerator.device, device=accelerator.device,
dtype=weight_dtype, dtype=weight_dtype,
) )
sigmas = expand_tensor_dims(sigmas, ndim=noise.ndim) sigmas = expand_tensor_dims(sigmas, ndim=noise.ndim)
@@ -797,38 +798,41 @@ class Trainer:
timesteps=timesteps, timesteps=timesteps,
) )
else: else:
if self.pose_condition: 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 # condition with the pose information
# normalize and patchify # normalize and patchify
# create noise for latent to feed diffusion transformer to get pred noise # create noise for latent to feed diffusion transformer to get pred noise
noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise target_video_latents = latent_conditions["latents"]
noisy_latents = noisy_latents.to(latent_conditions["latents"].dtype) 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 # add pose information to both channels
pose_noisy_latents = noisy_latents + pose_video_latents["latents"] pose_noisy_latents = noisy_latents + pose_video_latents["latents"]
pose_img_ref_latents = img_refs_latents["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 # 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"], num_frames=latent_conditions["num_frames"],
height=latent_conditions["height"], height=latent_conditions["height"],
width=latent_conditions["width"], width=latent_conditions["width"],
patch_size = 1, patch_size = 1,
patch_size_t = 1) patch_size_t = 1)
# target video as a input residual
noisy_latents_residual = post_conditioned_latent_patchify(latents=noisy_latents, noisy_latents_residual = post_conditioned_latent_patchify(latents=noisy_latents,
num_frames=latent_conditions["num_frames"], num_frames=latent_conditions["num_frames"],
height=latent_conditions["height"], height=latent_conditions["height"],
width=latent_conditions["width"], width=latent_conditions["width"],
patch_size = 1, patch_size = 1,
patch_size_t = 1) patch_size_t = 1)
# target latent patchified # Target Latent Patchified
# pose template noisey input [cat] img_ref + pose video
# pose information latent latent_conditions.update({"noisy_latents": pose_img_ref_latents["latents"]})
latent_conditions.update({"noisy_latents": latent_final["latents"]}) # input video noise at level residual information to adapter
latent_conditions.update({"noisy_latents_residual":noisy_latents_residual["latents"]}) latent_conditions.update({"noisy_latents_residual":noisy_latents_residual["latents"]})
else: else:
# Default to flow-matching noise addition # Default to flow-matching noise addition