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)
 
 

tests/test_dumas_h3_longvideos.py

@@ -1278,6 +1278,85 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
         self.assertEqual(chunk_length, 85)
         self.assertEqual(temporal_overlap, 17)
 
+    def test_unload_upscale_model_defers_while_held(self):
+        latent = importlib.import_module("dumas_h3_latent_upscale")
+
+        class FakeParam:
+            device = "cuda"
+
+        class FakeModel:
+            def __init__(self):
+                self.moves = []
+
+            def parameters(self):
+                return iter((FakeParam(),))
+
+            def to(self, device):
+                self.moves.append(device)
+                return self
+
+        cache_key = "upscale.safetensors::cuda::fp16"
+        original_cache_value = latent._MODEL_CACHE.get(cache_key)
+        original_hold_depth = latent._MODEL_HOLD_DEPTH
+        fake_model = FakeModel()
+        try:
+            latent._MODEL_CACHE[cache_key] = fake_model
+            latent._MODEL_HOLD_DEPTH = 0
+            with latent._hold_upscale_model_loaded():
+                latent.unload_upscale_model("upscale.safetensors", "cuda", "fp16")
+            self.assertEqual(fake_model.moves, [])
+
+            latent.unload_upscale_model("upscale.safetensors", "cuda", "fp16")
+            self.assertEqual(fake_model.moves, ["cpu"])
+        finally:
+            latent._MODEL_HOLD_DEPTH = original_hold_depth
+            if original_cache_value is None:
+                latent._MODEL_CACHE.pop(cache_key, None)
+            else:
+                latent._MODEL_CACHE[cache_key] = original_cache_value
+
+    def test_model_upscale_releases_cached_model_after_pass(self):
+        latent = importlib.import_module("dumas_h3_latent_upscale")
+        calls = []
+
+        original_temporal = latent._upscale_video_temporal_chunks
+        original_unload_now = latent._unload_upscale_model_now
+        original_cuda = latent.torch.cuda
+        original_device = getattr(latent.torch, "device", None)
+        try:
+            latent.torch.cuda = types.SimpleNamespace(is_available=lambda: True)
+            latent.torch.device = lambda value: value
+
+            def temporal(video, param, upscaler):
+                calls.append(("temporal", latent._MODEL_HOLD_DEPTH))
+                return "video", 8, 16
+
+            def unload_now(name, device, precision):
+                calls.append(("unload", name, device, precision, latent._MODEL_HOLD_DEPTH))
+
+            latent._upscale_video_temporal_chunks = temporal
+            latent._unload_upscale_model_now = unload_now
+
+            result = latent.upscale_latent_video("source", {
+                "mode": "model",
+                "model_name": "upscale.safetensors",
+                "device": "cuda",
+                "precision": "fp16",
+            })
+
+            self.assertEqual(result, ("video", 8, 16))
+            self.assertEqual(calls[0], ("temporal", 1))
+            self.assertEqual(calls[1], ("unload", "upscale.safetensors", "cuda", "fp16", 1))
+            self.assertEqual(latent._MODEL_HOLD_DEPTH, 0)
+        finally:
+            latent._upscale_video_temporal_chunks = original_temporal
+            latent._unload_upscale_model_now = original_unload_now
+            latent.torch.cuda = original_cuda
+            if original_device is None:
+                delattr(latent.torch, "device")
+            else:
+                latent.torch.device = original_device
+
     def test_upscale_video_model_raises_when_gpu_cannot_shrink(self):
         latent = importlib.import_module("dumas_h3_latent_upscale")