Bypass learned latent upscaler on low VRAM

2026-09-04 15:24:30 +00:00
parent 5dc5c2ce8b
commit a7f7225b3a
2 changed files with 146 additions and 5 deletions
@@ -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")