From 0b189dcf7a6489c169b3a5db45e4b341852e0ee6 Mon Sep 17 00:00:00 2001 From: Chris Dumas Date: Fri, 4 Sep 2026 11:04:26 +0000 Subject: [PATCH] Keep latent upscaler loaded during pass --- dumas_h3_latent_upscale.py | 30 +++++++++++- tests/test_dumas_h3_longvideos.py | 79 +++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 2 deletions(-) diff --git a/dumas_h3_latent_upscale.py b/dumas_h3_latent_upscale.py index 4e4df37..406e44d 100644 --- a/dumas_h3_latent_upscale.py +++ b/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) diff --git a/tests/test_dumas_h3_longvideos.py b/tests/test_dumas_h3_longvideos.py index 92fddb0..7522278 100644 --- a/tests/test_dumas_h3_longvideos.py +++ b/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")