Update nodes.py

This commit is contained in:
kijai
2024-02-07 20:59:43 +02:00
parent e7ddc001fa
commit 5987a54719
+16 -16
View File
@@ -337,7 +337,7 @@ class AD_MotionDirector_train:
validation_spatial_scale = 0.5
validation_seed = 44
validation_steps = 25
validation_steps = 50
validation_steps_tuple = [2, 25]
# Initialize distributed training
@@ -558,16 +558,16 @@ class AD_MotionDirector_train:
]
#Data batch sanity check
# if epoch == first_epoch and step == 0:
# "DO SANITY CHECK"
# do_sanity_check(
# pixel_values,
# cache_latents,
# validation_pipeline,
# device,
# output_dir=output_dir,
# text_prompt=text_prompt
# )
if epoch == first_epoch and step == 0:
"DO SANITY CHECK"
do_sanity_check(
pixel_values,
cache_latents,
validation_pipeline,
device,
output_dir=output_dir,
text_prompt=text_prompt
)
# Convert videos to latent space
@@ -630,15 +630,11 @@ class AD_MotionDirector_train:
if use_hflip:
model_pred_spatial = unet(noisy_latents_input, timesteps,
encoder_hidden_states=encoder_hidden_states).sample
model_pred_spatial.requires_grad_(True)
target_spatial.requires_grad_(True)
loss_spatial = F.mse_loss(model_pred_spatial[:, :, 0, :, :].float(),
target_spatial[:, :, 0, :, :].float(), reduction="mean")
else:
model_pred_spatial = unet(noisy_latents_input.unsqueeze(2), timesteps,
encoder_hidden_states=encoder_hidden_states).sample
model_pred_spatial.requires_grad_(True)
target_spatial.requires_grad_(True)
loss_spatial = F.mse_loss(model_pred_spatial[:, :, 0, :, :].float(),
target_spatial.float(), reduction="mean")
@@ -758,9 +754,10 @@ class AD_MotionDirector_train:
width = width,
).videos
save_videos_grid(sample, f"{output_dir}/samples/sample-{global_step}.gif")
print("samples shape ",samples.shape)
samples.append(sample)
unet.train()
unet.train()
samples = torch.concat(samples)
save_path = f"{output_dir}/samples/sample-{global_step}.gif"
@@ -787,6 +784,9 @@ class AD_MotionDirector_train:
import folder_paths
class DiffusersLoaderForTraining:
@classmethod
def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
return ""
@classmethod
def INPUT_TYPES(cls):
paths = []
for search_path in folder_paths.get_folder_paths("diffusers"):