From a7f7225b3af0747ee06648d4595ed233d6c4fd5e Mon Sep 17 00:00:00 2001 From: Chris Dumas Date: Fri, 4 Sep 2026 15:24:30 +0000 Subject: [PATCH] Bypass learned latent upscaler on low VRAM --- dumas_h3_latent_upscale.py | 54 +++++++++++++++-- tests/test_dumas_h3_longvideos.py | 97 +++++++++++++++++++++++++++++++ 2 files changed, 146 insertions(+), 5 deletions(-) diff --git a/dumas_h3_latent_upscale.py b/dumas_h3_latent_upscale.py index eb669e8..7295aa2 100644 --- a/dumas_h3_latent_upscale.py +++ b/dumas_h3_latent_upscale.py @@ -46,6 +46,7 @@ LOGGER = logging.getLogger(__name__) MP_UNIT = 1024 * 1024 RES_MULTIPLE = 32 CUDA_MODEL_CHUNK_LENGTH = 17 +CUDA_MODEL_MIN_TOTAL_VRAM_BYTES = 10 * 1024 * 1024 * 1024 _MODEL_HOLD_DEPTH = 0 @@ -76,6 +77,42 @@ def _should_retry_temporal_before_spatial(param): return _uses_cuda_model_upscale(param) and int(param.get("chunk_length", 0) or 0) > CUDA_MODEL_CHUNK_LENGTH +def _cuda_total_memory_bytes(): + try: + if not torch.cuda.is_available(): + return 0 + if hasattr(torch.cuda, "mem_get_info"): + _free, total = torch.cuda.mem_get_info() + return int(total) + current_device = torch.cuda.current_device() if hasattr(torch.cuda, "current_device") else 0 + props = torch.cuda.get_device_properties(current_device) + return int(getattr(props, "total_memory", 0) or 0) + except Exception: + return 0 + + +def _should_skip_cuda_model_upscale(param): + total = _cuda_total_memory_bytes() + return _uses_cuda_model_upscale(param) and 0 < total < CUDA_MODEL_MIN_TOTAL_VRAM_BYTES + + +def _fallback_to_interp(video, param, reason): + method = param.get("method", "bilinear") + LOGGER.warning("H3 latent upscale model %s; using %s interpolation instead", reason, method) + try: + gc.collect() + except Exception: + pass + if torch.cuda.is_available(): + try: + torch.cuda.empty_cache() + except Exception: + pass + interp_param = dict(param) + interp_param["mode"] = "interp" + return upscale_video_interp(video, interp_param) + + def _retry_with_smaller_temporal(video, param, upscaler, exc): if not _is_oom_error(exc): raise exc @@ -902,15 +939,22 @@ def upscale_latent_video(video, param): if mode == "off": return video, video.shape[-2], video.shape[-1] if mode == "model": + if _should_skip_cuda_model_upscale(param): + return _fallback_to_interp(video, param, "requires more than this card's VRAM") 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) + try: + 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) + except RuntimeError as exc: + if not _is_oom_error(exc): + raise + return _fallback_to_interp(video, param, "exhausted GPU memory") 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 cd75bca..b95d34f 100644 --- a/tests/test_dumas_h3_longvideos.py +++ b/tests/test_dumas_h3_longvideos.py @@ -1414,6 +1414,103 @@ class DumasH3LongVideosHelperTests(unittest.TestCase): else: latent.torch.device = original_device + def test_model_upscale_oom_falls_back_to_interp(self): + latent = importlib.import_module("dumas_h3_latent_upscale") + calls = [] + + original_temporal = latent._upscale_video_temporal_chunks + original_interp = latent.upscale_video_interp + 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, + empty_cache=lambda: calls.append(("empty_cache",)), + ) + latent.torch.device = lambda value: value + + def temporal(_video, _param, _upscaler): + calls.append(("temporal", latent._MODEL_HOLD_DEPTH)) + raise RuntimeError("H3 latent upscale exhausted its GPU spatial fallbacks") + + def interp(video, param): + calls.append(("interp", video, param.get("mode"), param.get("method"))) + return "interp_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.upscale_video_interp = interp + latent._unload_upscale_model_now = unload_now + + result = latent.upscale_latent_video("source", { + "mode": "model", + "model_name": "upscale.safetensors", + "method": "bilinear", + "device": "cuda", + "precision": "fp16", + }) + + self.assertEqual(result, ("interp_video", 8, 16)) + self.assertEqual(calls[0], ("temporal", 1)) + self.assertEqual(calls[1], ("unload", "upscale.safetensors", "cuda", "fp16", 1)) + self.assertIn(("empty_cache",), calls) + self.assertEqual(calls[-1], ("interp", "source", "interp", "bilinear")) + self.assertEqual(latent._MODEL_HOLD_DEPTH, 0) + finally: + latent._upscale_video_temporal_chunks = original_temporal + latent.upscale_video_interp = original_interp + 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_model_upscale_skips_learned_model_on_8gb_cuda(self): + latent = importlib.import_module("dumas_h3_latent_upscale") + calls = [] + + original_temporal = latent._upscale_video_temporal_chunks + original_interp = latent.upscale_video_interp + original_cuda = latent.torch.cuda + try: + latent.torch.cuda = types.SimpleNamespace( + is_available=lambda: True, + mem_get_info=lambda: (1 * 1024 * 1024 * 1024, 8 * 1024 * 1024 * 1024), + empty_cache=lambda: calls.append(("empty_cache",)), + ) + + def temporal(_video, _param, _upscaler): + calls.append(("temporal",)) + raise AssertionError("learned model path should be skipped on 8GB CUDA") + + def interp(video, param): + calls.append(("interp", video, param.get("mode"), param.get("method"))) + return "interp_video", 8, 16 + + latent._upscale_video_temporal_chunks = temporal + latent.upscale_video_interp = interp + + result = latent.upscale_latent_video("source", { + "mode": "model", + "model_name": "upscale.safetensors", + "method": "bilinear", + "device": "cuda", + "precision": "fp16", + }) + + self.assertEqual(result, ("interp_video", 8, 16)) + self.assertNotIn(("temporal",), calls) + self.assertIn(("empty_cache",), calls) + self.assertEqual(calls[-1], ("interp", "source", "interp", "bilinear")) + finally: + latent._upscale_video_temporal_chunks = original_temporal + latent.upscale_video_interp = original_interp + latent.torch.cuda = original_cuda + def test_upscale_video_model_raises_when_gpu_cannot_shrink(self): latent = importlib.import_module("dumas_h3_latent_upscale")