Post-refactor fixes

Author: semjonsona

README.md

@@ -21,10 +21,10 @@ video by [@graemeniedermayer](https://github.com/graemeniedermayer), more exampl
 images generated by [@semjon00](https://github.com/semjon00) from CC0 photos, more examples [here](https://github.com/thygate/stable-diffusion-webui-depthmap-script/pull/56#issuecomment-1367596463).
 
 ## Changelog
-* v0.3.13 
-    * Large code refactor
-    * Improved interface
-    * Slightly changed the behaviour of various options
+* v0.4.0 large code refactor
+    * UI improvements
+    * slightly changed the behaviour of various options
+    * extension may partially work even if some of the dependencies are unmet
 * v0.3.12
     * Fixed stereo image generation
     * Other bugfixes

scripts/core.py

@@ -230,7 +230,7 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
 
             if show_heat:
                 from dzoedepth.utils.misc import colorize
-                heatmap = colorize(img_output, cmap='inferno')
+                heatmap = Image.fromarray(colorize(img_output, cmap='inferno'))
                 generated_images[count]['heatmap'] = heatmap
 
             if gen_stereo:

scripts/depthmap_generation.py

@@ -3,6 +3,7 @@
 from PIL import Image
 from torchvision.transforms import Compose, transforms
 
+# TODO: depthmap_generation should not depend on WebUI
 from modules import shared, devices
 from modules.shared import opts, cmd_opts
 
@@ -29,7 +30,6 @@
 from pix2pix.options.test_options import TestOptions
 from pix2pix.models.pix2pix4depth_model import Pix2Pix4DepthModel
 
-
 # zoedepth
 from dzoedepth.models.builder import build_model
 from dzoedepth.utils.config import get_config
@@ -59,9 +59,6 @@ def ensure_models(self, model_type, device: torch.device, boost: bool):
 
     def load_models(self, model_type, device: torch.device, boost: bool):
         """Ensure that the depth model is loaded"""
-        # TODO: supply correct values for zoedepth
-        net_width = 512
-        net_height = 512
 
         # model path and name
         model_dir = "./models/midas"
@@ -171,22 +168,21 @@ def load_models(self, model_type, device: torch.device, boost: bool):
                 mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
             )
 
+        # When loading, zoedepth models will report the default net size.
+        # It will be overridden by the generation settings.
         elif model_type == 7:  # zoedepth_n
             print("zoedepth_n\n")
             conf = get_config("zoedepth", "infer")
-            conf.img_size = [net_width, net_height]
             model = build_model(conf)
 
         elif model_type == 8:  # zoedepth_k
             print("zoedepth_k\n")
             conf = get_config("zoedepth", "infer", config_version="kitti")
-            conf.img_size = [net_width, net_height]
             model = build_model(conf)
 
         elif model_type == 9:  # zoedepth_nk
             print("zoedepth_nk\n")
             conf = get_config("zoedepth_nk", "infer")
-            conf.img_size = [net_width, net_height]
             model = build_model(conf)
 
         model.eval()  # prepare for evaluation
@@ -221,15 +217,20 @@ def load_models(self, model_type, device: torch.device, boost: bool):
 
         devices.torch_gc()
 
-    def get_default_net_size(self, model_type):
+    @staticmethod
+    def get_default_net_size(model_type):
         # TODO: fill in, use in the GUI
         sizes = {
+            0: [448, 448],
             1: [512, 512],
             2: [384, 384],
             3: [384, 384],
             4: [384, 384],
             5: [384, 384],
             6: [256, 256],
+            7: [384, 512],
+            8: [384, 768],
+            9: [384, 512]
         }
         if model_type in sizes:
             return sizes[model_type]
@@ -254,8 +255,9 @@ def unload_models(self):
         self.device = None
 
     def get_raw_prediction(self, input, net_width, net_height):
-        """Get prediction from the model currently loaded by the class.
+        """Get prediction from the model currently loaded by the ModelHolder object.
         If boost is enabled, net_width and net_height will be ignored."""
+        # TODO: supply net size for zoedepth
         global device
         device = self.device
         # input image
@@ -264,17 +266,14 @@ def get_raw_prediction(self, input, net_width, net_height):
         if self.pix2pix_model is None:
             if self.depth_model_type == 0:
                 raw_prediction = estimateleres(img, self.depth_model, net_width, net_height)
-                raw_prediction_invert = True
             elif self.depth_model_type in [7, 8, 9]:
                 raw_prediction = estimatezoedepth(input, self.depth_model, net_width, net_height)
-                raw_prediction_invert = True
             else:
                 raw_prediction = estimatemidas(img, self.depth_model, net_width, net_height,
                                                self.resize_mode, self.normalization)
-                raw_prediction_invert = False
         else:
             raw_prediction = estimateboost(img, self.depth_model, self.depth_model_type, self.pix2pix_model)
-            raw_prediction_invert = False
+        raw_prediction_invert = self.depth_model_type in [0, 7, 8, 9]
         return raw_prediction, raw_prediction_invert
 
 

scripts/interface_webui.py

@@ -8,11 +8,12 @@
 from modules.shared import opts
 from modules.ui import plaintext_to_html
 from pathlib import Path
+from PIL import Image
 
 from scripts.gradio_args_transport import GradioComponentBundle
 from scripts.main import *
 from scripts.core import core_generation_funnel, unload_models, run_makevideo
-from PIL import Image
+from scripts.depthmap_generation import ModelHolder
 
 
 # Ugly workaround to fix gradio tempfile issue
@@ -102,19 +103,19 @@ def main_ui_panel(is_depth_tab):
         with gr.Group():
             with gr.Row():
                 inp += "gen_mesh", gr.Checkbox(
-                    label="Generate simple 3D mesh. "
-                          "(Fast, accurate only with ZoeDepth models and no boost, no custom maps)",
-                    value=False, visible=True)
+                    label="Generate simple 3D mesh", value=False, visible=True)
             with gr.Row(visible=False) as mesh_options_row_0:
+                gr.Label(value="Generates fast, accurate only with ZoeDepth models and no boost, no custom maps")
                 inp += "mesh_occlude", gr.Checkbox(label="Remove occluded edges", value=True, visible=True)
                 inp += "mesh_spherical", gr.Checkbox(label="Equirectangular projection", value=False, visible=True)
 
         if is_depth_tab:
             with gr.Group():
                 with gr.Row():
                     inp += "inpaint", gr.Checkbox(
-                        label="Generate 3D inpainted mesh. (Sloooow, required for generating videos)", value=False)
+                        label="Generate 3D inpainted mesh", value=False)
                 with gr.Group(visible=False) as inpaint_options_row_0:
+                    gr.Label("Generation is sloooow, required for generating videos")
                     inp += "inpaint_vids", gr.Checkbox(
                         label="Generate 4 demo videos with 3D inpainted mesh.", value=False)
                     gr.HTML("More options for generating video can be found in the Generate video tab")
@@ -139,6 +140,15 @@ def main_ui_panel(is_depth_tab):
 
         inp += "gen_normal", gr.Checkbox(label="Generate Normalmap (hidden! api only)", value=False, visible=False)
 
+        def update_delault_net_size(model_type):
+            w, h = ModelHolder.get_default_net_size(model_type)
+            return inp['net_width'].update(value=w), inp['net_height'].update(value=h)
+        inp['model_type'].change(
+            fn=update_delault_net_size,
+            inputs=inp['model_type'],
+            outputs=[inp['net_width'], inp['net_height']]
+        )
+
         inp['boost'].change(
             fn=lambda a, b: (options_depend_on_boost.update(visible=not a),
                              options_depend_on_match_size.update(visible=not a and not b)),

scripts/stereoimage_generation.py

@@ -1,4 +1,11 @@
-from numba import njit, prange
+try:
+    from numba import njit, prange
+except Exception as e:
+    print(f"WARINING! Numba failed to import! Stereoimage generation will be much slower! ({str(e)})")
+    from builtins import range as prange
+    def njit(parallel=False):
+        def Inner(func): return lambda *args, **kwargs: func(*args, **kwargs)
+        return Inner
 import numpy as np
 from PIL import Image
 
@@ -73,7 +80,7 @@ def apply_stereo_divergence(original_image, depth, divergence, separation, fill_
         )
 
 
-@njit
+@njit(parallel=False)
 def apply_stereo_divergence_naive(
         original_image, normalized_depth, divergence_px: float, separation_px: float, fill_technique):
     h, w, c = original_image.shape