From 5987a547196f0fd36ba365ec7a6f5fc7c31cec36 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 7 Feb 2024 20:59:43 +0200
Subject: [PATCH] Update nodes.py

---
 nodes.py | 32 ++++++++++++++++----------------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/nodes.py b/nodes.py
index 6cc4d1e..9b34362 100644
--- a/nodes.py
+++ b/nodes.py
@@ -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"):