Added residual

This commit is contained in:
CrossProduct
2025-01-22 20:10:01 +00:00
parent 981767154b
commit 731f326e86
2 changed files with 4 additions and 2 deletions
+3 -1
View File
@@ -220,10 +220,11 @@ def conditioned_forward_pass(
num_frames: int, num_frames: int,
height: int, height: int,
width: int, width: int,
noisy_latents_residual:torch.Tensor,
**kwargs, **kwargs,
) -> torch.Tensor: ) -> torch.Tensor:
rope_interpolation_scale = [1 / 25, 32, 32] rope_interpolation_scale = [1 / 25, 32, 32]
denoised_latents = transformer( denoised_latents = transformer(
hidden_states=noisy_latents, hidden_states=noisy_latents,
encoder_hidden_states=prompt_embeds, encoder_hidden_states=prompt_embeds,
@@ -234,6 +235,7 @@ def conditioned_forward_pass(
width=width, width=width,
rope_interpolation_scale=rope_interpolation_scale, rope_interpolation_scale=rope_interpolation_scale,
return_dict=False, return_dict=False,
residual_x=noisy_latents_residual,
)[0] )[0]
return {"latents": denoised_latents} return {"latents": denoised_latents}
+1 -1
View File
@@ -827,6 +827,7 @@ class Trainer:
patch_size_t = 1) patch_size_t = 1)
# pose information latent # pose information latent
latent_conditions.update({"noisy_latents": latent_final["latents"]}) latent_conditions.update({"noisy_latents": latent_final["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
noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise
@@ -848,7 +849,6 @@ class Trainer:
transformer=self.transformer, transformer=self.transformer,
scheduler=self.scheduler, scheduler=self.scheduler,
timesteps=timesteps, timesteps=timesteps,
residual_x=noisy_latents_residual,
**latent_conditions, **latent_conditions,
**text_conditions, **text_conditions,
) )