Keep latent upscaler loaded during pass

dumas_h3_latent_upscale.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from functools import lru_cache
 import gc
 import glob
@@ -43,6 +44,7 @@ _LATENT_UPSCALE_FOLDER = "latent_upscale_models"
 MP_UNIT = 1024 * 1024
 RES_MULTIPLE = 32
 CUDA_MODEL_CHUNK_LENGTH = 17
+_MODEL_HOLD_DEPTH = 0
 
 
 def _uses_cuda_model_upscale(param):
@@ -349,7 +351,7 @@ def load_upscale_model(name, device, precision):
     return model
 
 
-def unload_upscale_model(name, device, precision):
+def _unload_upscale_model_now(name, device, precision):
     cache_key = f"{name}::{device}::{precision}"
     model = _MODEL_CACHE.get(cache_key)
     if model is not None and str(next(model.parameters()).device) != "cpu":
@@ -361,6 +363,22 @@ def unload_upscale_model(name, device, precision):
             pass
 
 
+def unload_upscale_model(name, device, precision):
+    if _MODEL_HOLD_DEPTH > 0:
+        return
+    _unload_upscale_model_now(name, device, precision)
+
+
+@contextmanager
+def _hold_upscale_model_loaded():
+    global _MODEL_HOLD_DEPTH
+    _MODEL_HOLD_DEPTH += 1
+    try:
+        yield
+    finally:
+        _MODEL_HOLD_DEPTH -= 1
+
+
 def _compute_upscale_target(width, height, h_in, w_in):
     ds = 16
     w_px = float(width)
@@ -819,7 +837,15 @@ def upscale_latent_video(video, param):
     if mode == "off":
         return video, video.shape[-2], video.shape[-1]
     if mode == "model":
-        return _upscale_video_temporal_chunks(video, param, upscale_video_model)
+        model_name = param.get("model_name")
+        device = param.get("device", "cuda")
+        precision = param.get("precision", "fp16")
+        dev = torch.device(device if (device == "cpu" or torch.cuda.is_available()) else "cpu")
+        with _hold_upscale_model_loaded():
+            try:
+                return _upscale_video_temporal_chunks(video, param, upscale_video_model)
+            finally:
+                _unload_upscale_model_now(model_name, dev, precision)
     return _upscale_video_temporal_chunks(video, param, upscale_video_interp)