initial commit

extensions-builtin/forge_preprocessor_tile/scripts/preprocessor_tile.py

@@ -0,0 +1,102 @@
+import torch
+
+from modules_forge.supported_preprocessor import Preprocessor, PreprocessorParameter
+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__()
+        self.name = 'tile_resample'
+        self.tags = ['Tile']
+        self.model_filename_filters = ['tile']
+        self.slider_resolution = PreprocessorParameter(visible=False)
+        self.latent = None
+
+    def register_latent(self, process, cond):
+        vae = process.sd_model.forge_objects.vae
+        # This is a powerful VAE with integrated memory management, bf16, and tiled fallback.
+
+        latent_image = vae.encode(cond.movedim(1, -1))
+        latent_image = process.sd_model.forge_objects.vae.first_stage_model.process_in(latent_image)
+        self.latent = latent_image
+        return self.latent
+
+
+class PreprocessorTileColorFix(PreprocessorTile):
+    def __init__(self):
+        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, mask, *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.predictor.sigma_data
+
+        if getattr(process, 'is_hr_pass', False):
+            k = int(self.variation * 2)
+        else:
+            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
+
+        return cond, mask
+
+
+class PreprocessorTileColorFixSharp(PreprocessorTileColorFix):
+    def __init__(self):
+        super().__init__()
+        self.name = 'tile_colorfix+sharp'
+        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, mask, *args, **kwargs):
+        self.sharpness = float(kwargs['unit'].threshold_b)
+        return super().process_before_every_sampling(process, cond, mask, *args, **kwargs)
+
+
+add_supported_preprocessor(PreprocessorTile())
+
+add_supported_preprocessor(PreprocessorTileColorFix())
+
+add_supported_preprocessor(PreprocessorTileColorFixSharp())