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

Author: semjonsona

scripts/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()

scripts/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:

scripts/interface_webui.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():