Tiling mode

This is very cool, but it still has some seams for some reason ;(
Drawbacks: Clutters UI

Closes #444

Author: semjon00

src/common_constants.py

@@ -21,6 +21,7 @@ def __init__(self, default_value=None, *args):
     NET_SIZE_MATCH = False
     NET_WIDTH = 448
     NET_HEIGHT = 448
+    TILING_MODE = False
 
     DO_OUTPUT_DEPTH = True
     OUTPUT_DEPTH_INVERT = False

src/common_ui.py

@@ -48,6 +48,10 @@ def main_ui_panel(is_depth_tab):
             with gr.Row(visible=False) as options_depend_on_match_size:
                 inp += go.NET_WIDTH, gr.Slider(minimum=64, maximum=2048, step=64, label='Net width')
                 inp += go.NET_HEIGHT, gr.Slider(minimum=64, maximum=2048, step=64, label='Net height')
+            with gr.Row():
+                inp += go.TILING_MODE, gr.Checkbox(
+                    label='Tiling mode', info='Reduces seams that appear if the depthmap is tiled into a grid'
+                )
 
         with gr.Box() as cur_option_root:
             inp -= 'depthmap_gen_row_2', cur_option_root
@@ -75,7 +79,7 @@ def main_ui_panel(is_depth_tab):
 
         with gr.Box():
             with gr.Row():
-                inp += go.GEN_STEREO, gr.Checkbox(label="Generate stereoscopic image(s)")
+                inp += go.GEN_STEREO, gr.Checkbox(label="Generate stereoscopic (3D) image(s)")
             with gr.Column(visible=False) as stereo_options:
                 with gr.Row():
                     inp += go.STEREO_MODES, gr.CheckboxGroup(
@@ -178,6 +182,13 @@ def update_default_net_size(model_type):
             outputs=[inp[go.NET_SIZE_MATCH], options_depend_on_match_size]
         )
         inp.add_rule(options_depend_on_match_size, 'visible-if-not', go.NET_SIZE_MATCH)
+        inp[go.TILING_MODE].change(  # Go boost! Wroom!..
+            fn=lambda a: (
+                inp[go.BOOST].update(value=False), inp[go.NET_SIZE_MATCH].update(value=True)
+            ) if a else (inp[go.BOOST].update(), inp[go.NET_SIZE_MATCH].update()),
+            inputs=[inp[go.TILING_MODE]],
+            outputs=[inp[go.BOOST], inp[go.NET_SIZE_MATCH]]
+        )
 
         inp.add_rule(options_depend_on_output_depth_1, 'visible-if', go.DO_OUTPUT_DEPTH)
         inp.add_rule(go.OUTPUT_DEPTH_INVERT, 'visible-if', go.DO_OUTPUT_DEPTH)

src/core.py

@@ -121,7 +121,7 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
     try:
         if not inputdepthmaps_complete:
             print("Loading model(s) ..")
-            model_holder.ensure_models(inp[go.MODEL_TYPE], device, inp[go.BOOST])
+            model_holder.ensure_models(inp[go.MODEL_TYPE], device, inp[go.BOOST], inp[go.TILING_MODE])
         print("Computing output(s) ..")
         # iterate over input images
         for count in trange(0, len(inputimages)):
@@ -170,6 +170,7 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
                 # override net size (size may be different for different images)
                 if inp[go.NET_SIZE_MATCH]:
                     # Round up to a multiple of 32 to avoid potential issues
+                    # TODO: buggs for Depth Anything
                     net_width = (inputimages[count].width + 31) // 32 * 32
                     net_height = (inputimages[count].height + 31) // 32 * 32
                 else:

src/depthmap_generation.py

@@ -43,6 +43,7 @@ def __init__(self):
         # Extra stuff
         self.resize_mode = None
         self.normalization = None
+        self.tiling_mode = False
 
 
     def update_settings(self, **kvargs):
@@ -51,18 +52,23 @@ def update_settings(self, **kvargs):
             setattr(self, k, v)
 
 
-    def ensure_models(self, model_type, device: torch.device, boost: bool):
+    def ensure_models(self, model_type, device: torch.device, boost: bool, tiling_mode: bool = False):
         # TODO: could make it more granular
         if model_type == -1 or model_type is None:
             self.unload_models()
             return
         # Certain optimisations are irreversible and not device-agnostic, thus changing device requires reloading
-        if model_type != self.depth_model_type or boost != (self.pix2pix_model is not None) or device != self.device:
+        if (
+                model_type != self.depth_model_type or
+                boost != (self.pix2pix_model is not None) or
+                device != self.device or
+                tiling_mode != self.tiling_mode
+        ):
             self.unload_models()
-            self.load_models(model_type, device, boost)
+            self.load_models(model_type, device, boost, tiling_mode)
         self.reload()
 
-    def load_models(self, model_type, device: torch.device, boost: bool):
+    def load_models(self, model_type, device: torch.device, boost: bool, tiling_mode: bool = False):
         """Ensure that the depth model is loaded"""
 
         # TODO: we need to at least try to find models downloaded by other plugins (e.g. controlnet)
@@ -205,7 +211,6 @@ def load_models(self, model_type, device: torch.device, boost: bool):
                 model.enable_xformers_memory_efficient_attention()
             except:
                 pass  # run without xformers
-
         elif model_type == 11:  # depth_anything
             from depth_anything.dpt import DPT_DINOv2
             # This will download the model... to some place
@@ -223,6 +228,17 @@ def load_models(self, model_type, device: torch.device, boost: bool):
 
             model.load_state_dict(torch.load(model_path))
 
+        if tiling_mode:
+            def flatten(el):
+                flattened = [flatten(children) for children in el.children()]
+                res = [el]
+                for c in flattened:
+                    res += c
+                return res
+            layers = flatten(model)  # Hijacking the model
+            for layer in [layer for layer in layers if type(layer) == torch.nn.Conv2d or type(layer) == torch.nn.Conv1d]:
+                layer.padding_mode = 'circular'
+
         if model_type in range(0, 10):
             model.eval()  # prepare for evaluation
         # optimize
@@ -238,6 +254,7 @@ def load_models(self, model_type, device: torch.device, boost: bool):
         self.depth_model_type = model_type
         self.resize_mode = resize_mode
         self.normalization = normalization
+        self.tiling_mode = tiling_mode
 
         self.device = device