Post-refactor fixes vol. 2

* Reload model before generation, if it is offloaded to CPU
* Load model if boost got selected
* Do not try to offload pix2pix
* Net dimensions are multiple of 32 regardless of match size
* Change the default net size to default net size of the default model
* Fixed script mode
* UI fixes

Author: semjon00

scripts/depthmap.py

@@ -49,8 +49,8 @@ def main_ui_panel(is_depth_tab):
                 with gr.Group(visible=False) as options_depend_on_boost:
                     inp += 'match_size', gr.Checkbox(label="Match net size to input size", value=False)
             with gr.Row(visible=False) as options_depend_on_match_size:
-                inp += 'net_width', gr.Slider(minimum=64, maximum=2048, step=64, label='Net width', value=512)
-                inp += 'net_height', gr.Slider(minimum=64, maximum=2048, step=64, label='Net height', value=512)
+                inp += 'net_width', gr.Slider(minimum=64, maximum=2048, step=64, label='Net width', value=448)
+                inp += 'net_height', gr.Slider(minimum=64, maximum=2048, step=64, label='Net height', value=448)
 
         with gr.Group():
             with gr.Row():
@@ -104,18 +104,20 @@ def main_ui_panel(is_depth_tab):
             with gr.Row():
                 inp += "gen_mesh", gr.Checkbox(
                     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)
+            with gr.Group(visible=False) as mesh_options:
+                with gr.Row():
+                    gr.HTML(value="Generates fast, accurate only with ZoeDepth models and no boost, no custom maps")
+                with gr.Row():
+                    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", value=False)
                 with gr.Group(visible=False) as inpaint_options_row_0:
-                    gr.Label("Generation is sloooow, required for generating videos")
+                    gr.HTML("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")
@@ -199,9 +201,9 @@ def stereo_options_visibility(v):
         )
 
         inp['gen_mesh'].change(
-            fn=lambda v: mesh_options_row_0.update(visible=v),
+            fn=lambda v: mesh_options.update(visible=v),
             inputs=[inp['gen_mesh']],
-            outputs=[mesh_options_row_0]
+            outputs=[mesh_options]
         )
 
         def inpaint_options_visibility(v):

src/core.py

@@ -62,7 +62,7 @@ def reload_sd_model():
 def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp):
     if len(inputimages) == 0 or inputimages[0] is None:
         return [], '', ''
-    if len(inputdepthmaps) == 0:
+    if inputdepthmaps is None or len(inputdepthmaps) == 0:
         inputdepthmaps: list[Image] = [None for _ in range(len(inputimages))]
     inputdepthmaps_complete = all([x is not None for x in inputdepthmaps])
 
@@ -78,8 +78,8 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
     gen_mesh = inp["gen_mesh"]
     gen_normal = inp["gen_normal"] if "gen_normal" in inp else False
     gen_stereo = inp["gen_stereo"]
-    inpaint = inp["inpaint"]
-    inpaint_vids = inp["inpaint_vids"]
+    inpaint = inp["inpaint"] if "inpaint" in inp else False
+    inpaint_vids = inp["inpaint_vids"] if "inpaint_vids" in inp else False
     invert_depth = inp["invert_depth"]
     match_size = inp["match_size"]
     mesh_occlude = inp["mesh_occlude"]
@@ -165,7 +165,9 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
             else:
                 # override net size (size may be different for different images)
                 if match_size:
-                    net_width, net_height = inputimages[count].width, inputimages[count].height
+                    # Round up to a multiple of 32 to avoid potential issues
+                    net_width = (inputimages[count].width + 31) // 32 * 32
+                    net_height = (inputimages[count].height + 31) // 32 * 32
                 raw_prediction, raw_prediction_invert = \
                     model_holder.get_raw_prediction(inputimages[count], net_width, net_height)
 
@@ -304,14 +306,14 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
         else:
             raise e
     finally:
-        if not (hasattr(opts, 'depthmap_script_keepmodels') and opts.depthmap_script_keepmodels):
+        if hasattr(opts, 'depthmap_script_keepmodels') and opts.depthmap_script_keepmodels:
+            model_holder.offload()  # Swap to CPU memory
+        else:
             if 'model' in locals():
                 del model
             if 'pix2pixmodel' in locals():
                 del pix2pix_model
             model_holder.unload_models()
-        else:
-            model_holder.swap_to_cpu_memory()
 
         gc.collect()
         devices.torch_gc()

src/depthmap_generation.py

@@ -42,6 +42,7 @@ def __init__(self):
         self.pix2pix_model = None
         self.depth_model_type = None
         self.device = None  # Target device, the model may be swapped from VRAM into RAM.
+        self.offloaded = False  # True means current device is not the target device
 
         # Extra stuff
         self.resize_mode = None
@@ -53,9 +54,10 @@ def ensure_models(self, model_type, device: torch.device, boost: bool):
             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:
             self.unload_models()
             self.load_models(model_type, device, boost)
+        self.reload()
 
     def load_models(self, model_type, device: torch.device, boost: bool):
         """Ensure that the depth model is loaded"""
@@ -236,11 +238,24 @@ def get_default_net_size(model_type):
             return sizes[model_type]
         return [512, 512]
 
-    def swap_to_cpu_memory(self):
+    def offload(self):
+        """Move to RAM to conserve VRAM"""
+        if self.device != torch.device('cpu') and not self.offloaded:
+            self.move_models_to(torch.device('cpu'))
+            self.offloaded = True
+
+    def reload(self):
+        """Undoes offload"""
+        if self.offloaded:
+            self.move_models_to(self.device)
+            self.offloaded = True
+
+    def move_models_to(self, device):
         if self.depth_model is not None:
-            self.depth_model.to(torch.device('cpu'))
+            self.depth_model.to(device)
         if self.pix2pix_model is not None:
-            self.pix2pix_model.to(torch.device('cpu'))
+            pass
+            # TODO: pix2pix offloading not implemented
 
     def unload_models(self):
         if self.depth_model is not None or self.pix2pix_model is not None: