Post-refactor fixes vol. 2

* Reload model before generation, if it is offloaded to CPU
* Load different model if boost got selected

Author: semjonsona

scripts/core.py

@@ -304,14 +304,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,23 @@ 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'))
+            self.pix2pix_model.to(device)
 
     def unload_models(self):
         if self.depth_model is not None or self.pix2pix_model is not None: