From 5775fedb7b43f9d50200e6cc72a6a7e93e221ec8 Mon Sep 17 00:00:00 2001
From: PLEXATIC <plexatic@gmail.com>
Date: Sat, 6 Apr 2024 17:41:32 +0200
Subject: [PATCH] Added logic to skip first n layers per block

The paper states that the process should only be applied for the layers 24 and later
---
 nodes.py | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index 88855d5..c564761 100644
--- a/nodes.py
+++ b/nodes.py
@@ -64,7 +64,11 @@ class ApplyVisualStyle:
         input_blocks,
         middle_block,
         output_blocks,
-        init_image = None
+        init_image = None,
+        skip_input_layers=0,
+        skip_middle_layers=0,
+        skip_output_layers=0
+
     ):
         positive = copy_cond(positive)
         reference_cond = copy_cond(reference_cond)
@@ -74,10 +78,27 @@ class ApplyVisualStyle:
 
         block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks)
 
+
+        layer_indexes = {
+            "input": 0,
+            "middle": 0,
+            "output": 0
+        }
+
+        n_skip_per_block = {
+            "input": skip_input_layers,
+            "middle": skip_middle_layers,
+            "output": skip_output_layers
+        }
+
         for n, m in model.model.diffusion_model.named_modules():
             if m.__class__.__name__  == "CrossAttention":
                 is_enabled = self.activate_block_choice(n, block_choices)
-
+                if is_enabled:
+                    block_name = n.split("_")[0]
+                    if layer_indexes[block_name] < n_skip_per_block[block_name]:
+                        is_enabled = False
+                    layer_indexes[block_name] += 1
                 if hasattr(m.forward, 'module_self'):
                     m.forward.enabled = is_enabled and enabled
                 else: