storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-09 00:09:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

reformulate the conditioning and patchification

Browse Source
This commit is contained in:
CrossProduct
2025-01-18 03:11:09 +00:00
parent 59afd55775
commit 44df9e3bce
2 changed files with 46 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/conditioning/condition_latents_prepare.py
+11 -5
View File
@@ -21,8 +21,9 @@ def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int
post_patch_height = height // patch_size
post_patch_width = width // patch_size
dim1 = num_frames // patch_size_t * height // patch_size * width // patch_size
dim2 = num_channels * patch_size_t * patch_size * patch_size
# to check math lol
# dim1 = num_frames // patch_size_t * height // patch_size * width // patch_size
# dim2 = num_channels * patch_size_t * patch_size * patch_size
latents = latents.reshape(
batch_size,
@@ -75,6 +76,11 @@ 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 = _pack_latents(latents, patch_size, patch_size_t)
return {"latents": latents, "num_frames": num_frames, "height": height, "width": width}
# 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,
"height": height,
"width": width,
"mean":vae.latents_mean,
"std":vae.latents_std}
finetrainers/trainer.py
+35 -17
View File
@@ -748,6 +748,8 @@ class Trainer:
)
img_refs_latents = make_contiguous(img_refs_latents)
text_conditions = make_contiguous(text_conditions)
@@ -773,23 +775,13 @@ class Trainer:
generator=self.state.generator,
)
timesteps = (sigmas * 1000.0).long()
if self.pose_condition:
noise = torch.randn(
noise = torch.randn(
latent_conditions["latents"].shape,
generator=self.state.generator,
device=accelerator.device,
dtype=weight_dtype,
)
# B x 2C latent
imf_ref_noise = torch.cat([img_refs_latents,noise])
# create noise for the img ref and concat it to latent to be shape B, 2c, etc.
else:
noise = torch.randn(
latent_conditions["latents"].shape,
generator=self.state.generator,
device=accelerator.device,
dtype=weight_dtype,
)
)
sigmas = expand_tensor_dims(sigmas, ndim=noise.ndim)
@@ -803,13 +795,39 @@ class Trainer:
timesteps=timesteps,
)
else:
if self.pose_condition:
# Default to flow-matching noise addition
# 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
noisy_latents = noisy_latents.to(latent_conditions["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)
# 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)
# this is used for whatever reason to pass into the
latent_conditions.update({"noisy_latents": noisy_latents})
latent_conditions.update({"noisy_latents": latent_final})
else:
# Default to flow-matching noise addition
noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise
@@ -831,7 +849,7 @@ class Trainer:
transformer=self.transformer,
scheduler=self.scheduler,
timesteps=timesteps,
**img_refs_latents,
**latent_conditions,
**text_conditions,
)