From 65bb421d8db29acae156b667dc8d8a2a6aaabe68 Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Wed, 8 Nov 2023 15:03:59 -0600
Subject: [PATCH] Brought motion scale masking feature to the experimental
 nodes, too

---
 animatediff/nodes_experimental.py | 40 ++++++++++++++++++++++++-------
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/animatediff/nodes_experimental.py b/animatediff/nodes_experimental.py
index 0e31347..15c41ea 100644
--- a/animatediff/nodes_experimental.py
+++ b/animatediff/nodes_experimental.py
@@ -12,9 +12,9 @@ class AnimateDiffModelSettingsSimple:
                 "motion_pe_stretch": ("INT", {"default": 0, "min": 0, "step": 1}),
             },
             "optional": {
-                "mask": ("MASK",),
-                "min_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
-                "max_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
+                "mask_motion_scale": ("MASK",),
+                "min_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
+                "max_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
             }
         }
     
@@ -22,9 +22,13 @@ class AnimateDiffModelSettingsSimple:
     CATEGORY = "Animate Diff 🎭🅐🅓/motion settings/experimental"
     FUNCTION = "get_motion_model_settings"
 
-    def get_motion_model_settings(self, motion_pe_stretch: int, mask: torch.Tensor=None, min_scale: float=1.0, max_scale: float=1.0):
+    def get_motion_model_settings(self, motion_pe_stretch: int,
+                                  mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
         motion_model_settings = MotionModelSettings(
-            motion_pe_stretch=motion_pe_stretch
+            motion_pe_stretch=motion_pe_stretch,
+            mask_attn_scale=mask_motion_scale,
+            mask_attn_scale_min=min_motion_scale,
+            mask_attn_scale_max=max_motion_scale,
             )
 
         return (motion_model_settings,)
@@ -44,6 +48,11 @@ class AnimateDiffModelSettingsAdvanced:
                 "initial_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
                 "final_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
             },
+            "optional": {
+                "mask_motion_scale": ("MASK",),
+                "min_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
+                "max_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
+            }
         }
     
     RETURN_TYPES = ("MOTION_MODEL_SETTINGS",)
@@ -53,7 +62,8 @@ class AnimateDiffModelSettingsAdvanced:
     def get_motion_model_settings(self, pe_strength: float, attn_strength: float, other_strength: float,
                                   motion_pe_stretch: int,
                                   cap_initial_pe_length: int, interpolate_pe_to_length: int,
-                                  initial_pe_idx_offset: int, final_pe_idx_offset: int):
+                                  initial_pe_idx_offset: int, final_pe_idx_offset: int,
+                                  mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
         motion_model_settings = MotionModelSettings(
             pe_strength=pe_strength,
             attn_strength=attn_strength,
@@ -62,7 +72,10 @@ class AnimateDiffModelSettingsAdvanced:
             interpolate_pe_to_length=interpolate_pe_to_length,
             initial_pe_idx_offset=initial_pe_idx_offset,
             final_pe_idx_offset=final_pe_idx_offset,
-            motion_pe_stretch=motion_pe_stretch
+            motion_pe_stretch=motion_pe_stretch,
+            mask_attn_scale=mask_motion_scale,
+            mask_attn_scale_min=min_motion_scale,
+            mask_attn_scale_max=max_motion_scale,
             )
 
         return (motion_model_settings,)
@@ -87,6 +100,11 @@ class AnimateDiffModelSettingsAdvancedAttnStrengths:
                 "initial_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
                 "final_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
             },
+            "optional": {
+                "mask_motion_scale": ("MASK",),
+                "min_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
+                "max_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
+            }
         }
     
     RETURN_TYPES = ("MOTION_MODEL_SETTINGS",)
@@ -102,7 +120,8 @@ class AnimateDiffModelSettingsAdvancedAttnStrengths:
                                   other_strength: float,
                                   motion_pe_stretch: int,
                                   cap_initial_pe_length: int, interpolate_pe_to_length: int,
-                                  initial_pe_idx_offset: int, final_pe_idx_offset: int):
+                                  initial_pe_idx_offset: int, final_pe_idx_offset: int,
+                                  mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
         motion_model_settings = MotionModelSettings(
             pe_strength=pe_strength,
             attn_strength=attn_strength,
@@ -116,7 +135,10 @@ class AnimateDiffModelSettingsAdvancedAttnStrengths:
             interpolate_pe_to_length=interpolate_pe_to_length,
             initial_pe_idx_offset=initial_pe_idx_offset,
             final_pe_idx_offset=final_pe_idx_offset,
-            motion_pe_stretch=motion_pe_stretch
+            motion_pe_stretch=motion_pe_stretch,
+            mask_attn_scale=mask_motion_scale,
+            mask_attn_scale_min=min_motion_scale,
+            mask_attn_scale_max=max_motion_scale,
             )
 
         return (motion_model_settings,)