Bypass learned latent upscaler on low VRAM

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)
 
 

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")