Retested changes without conditioning it works.

This commit is contained in:
CrossProduct
2025-01-17 22:43:20 +00:00
parent b73d61ae4d
commit 59afd55775
+6 -6
View File
@@ -100,10 +100,10 @@ class Trainer:
self.state.model_name = self.args.model_name self.state.model_name = self.args.model_name
self.model_config = get_config_from_model_name(self.args.model_name, self.args.training_type) self.model_config = get_config_from_model_name(self.args.model_name, self.args.training_type)
self.pose_condition = False
if self.args.pose_column != None: if self.args.pose_column != None:
self.pose_condition = True self.pose_condition = True
else:
self.pose_condition = False
def prepare_dataset(self) -> None: def prepare_dataset(self) -> None:
# TODO(aryan): Make a background process for fetching # TODO(aryan): Make a background process for fetching
@@ -723,7 +723,7 @@ class Trainer:
latent_conditions = make_contiguous(latent_conditions) latent_conditions = make_contiguous(latent_conditions)
if self.pose_conditioning: if self.pose_condition:
# VAE output not patchified yet. # VAE output not patchified yet.
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,
@@ -739,7 +739,7 @@ class Trainer:
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_latents, image_or_video=img_refs,
patch_size=self.transformer_config.patch_size, patch_size=self.transformer_config.patch_size,
patch_size_t=self.transformer_config.patch_size_t, patch_size_t=self.transformer_config.patch_size_t,
device=accelerator.device, device=accelerator.device,
@@ -773,7 +773,7 @@ class Trainer:
generator=self.state.generator, generator=self.state.generator,
) )
timesteps = (sigmas * 1000.0).long() timesteps = (sigmas * 1000.0).long()
if self.pose_conditioning: if self.pose_condition:
noise = torch.randn( noise = torch.randn(
latent_conditions["latents"].shape, latent_conditions["latents"].shape,
generator=self.state.generator, generator=self.state.generator,
@@ -826,7 +826,7 @@ class Trainer:
) )
weights = expand_tensor_dims(weights, noise.ndim) weights = expand_tensor_dims(weights, noise.ndim)
if self.pose_conditioning: if self.pose_condition:
pred = self.model_config["forward_pass"]( pred = self.model_config["forward_pass"](
transformer=self.transformer, transformer=self.transformer,
scheduler=self.scheduler, scheduler=self.scheduler,