Update preprocessor_tile.py

extensions-builtin/forge_preprocessor_tile/scripts/preprocessor_tile.py

@@ -1,10 +1,15 @@
 import torch
 
 from modules_forge.supported_preprocessor import Preprocessor, PreprocessorParameter
-from modules_forge.forge_util import numpy_to_pytorch, resize_image_with_pad
 from modules_forge.shared import add_supported_preprocessor
 
 
+def blur(x, k):
+    y = torch.nn.functional.pad(x, (k, k, k, k), mode='replicate')
+    y = torch.nn.functional.avg_pool2d(y, (k * 2 + 1, k * 2 + 1), stride=(1, 1))
+    return y
+
+
 class PreprocessorTile(Preprocessor):
     def __init__(self):
         super().__init__()
@@ -29,11 +34,46 @@ class PreprocessorTileColorFix(PreprocessorTile):
         super().__init__()
         self.name = 'tile_colorfix'
         self.slider_1 = PreprocessorParameter(label='Variation', value=8.0, minimum=3.0, maximum=32.0, step=1.0, visible=True)
+        self.variation = 8
+        self.sharpness = None
 
     def process_before_every_sampling(self, process, cond, *args, **kwargs):
+        self.variation = int(kwargs['unit'].threshold_a)
+
         latent = self.register_latent(process, cond)
 
         unet = process.sd_model.forge_objects.unet.clone()
+        sigma_data = process.sd_model.forge_objects.unet.model.model_sampling.sigma_data
+        k = int(self.variation)
+
+        def block_proc(h, flag, transformer_options):
+            location, block_id = transformer_options['block']
+            cond_mark = transformer_options['cond_mark'][:, None, None, None]  # cond is 0
+
+            if location == 'input' and block_id == 0 and flag == 'before':
+                sigma = transformer_options['sigmas'].to(h)
+                self.x_input = h[:, :4]  # Inpaint fix
+                self.x_input_sigma_space = self.x_input * (sigma ** 2 + sigma_data ** 2) ** 0.5
+
+            if location == 'last' and block_id == 0 and flag == 'after':
+                sigma = transformer_options['sigmas'].to(h)
+                eps_estimation = h[:, :4]
+                denoised = self.x_input_sigma_space - eps_estimation * sigma
+
+                denoised = denoised - blur(denoised, k) + blur(latent.to(denoised), k)
+
+                if isinstance(self.sharpness, float):
+                    detail_weight = float(self.sharpness) * 0.01
+                    neg = detail_weight * blur(denoised, k) + (1 - detail_weight) * denoised
+                    denoised = (1 - cond_mark) * denoised + cond_mark * neg
+
+                eps_modified = (self.x_input_sigma_space - denoised) / sigma
+
+                return eps_modified
+
+            return h
+
+        unet.add_block_modifier(block_proc)
 
         process.sd_model.forge_objects.unet = unet
 
@@ -47,12 +87,8 @@ class PreprocessorTileColorFixSharp(PreprocessorTileColorFix):
         self.slider_2 = PreprocessorParameter(label='Sharpness', value=1.0, minimum=0.0, maximum=2.0, step=0.01, visible=True)
 
     def process_before_every_sampling(self, process, cond, *args, **kwargs):
-        latent = self.register_latent(process, cond)
-
-        unet = process.sd_model.forge_objects.unet.clone()
-
-        process.sd_model.forge_objects.unet = unet
-
+        self.sharpness = float(kwargs['unit'].threshold_b)
+        super().process_before_every_sampling(process, cond, *args, **kwargs)
         return