mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-18 00:33:52 +00:00
removed some code made things clearer to replace the pipeline inference which is needed for validation.
This commit is contained in:
@@ -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}
|
|
||||||
@@ -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,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user