Prioritize temporal fallback for latent upscale OOM

2026-09-04 15:10:38 +00:00
parent 83645811f4
commit 5dc5c2ce8b
2 changed files with 43 additions and 0 deletions
@@ -1290,6 +1290,40 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
        self.assertEqual(chunk_length, 85)
        self.assertEqual(temporal_overlap, 17)

+    def test_cuda_model_oom_retries_temporal_before_spatial_fallback(self):
+        latent = importlib.import_module("dumas_h3_latent_upscale")
+
+        calls = []
+        original_tiled = latent._upscale_video_model_tiled
+        original_shrink_model = latent._shrink_model_tile_param
+        try:
+            def tiled(_video, param):
+                calls.append(("tiled", param.get("chunk_length"), param.get("tile_size_mode")))
+                raise RuntimeError("out of memory")
+
+            latent._upscale_video_model_tiled = tiled
+            latent._shrink_model_tile_param = (
+                lambda param: calls.append(("shrink_spatial", param.get("tile_size_mode"))) or None
+            )
+
+            with self.assertRaisesRegex(RuntimeError, "smaller temporal chunk"):
+                latent.upscale_video_model(
+                    "video",
+                    {
+                        "mode": "model",
+                        "device": "cuda",
+                        "model_name": "upscale.safetensors",
+                        "chunk_length": 85,
+                        "temporal_overlap": 17,
+                    },
+                )
+
+            self.assertEqual(calls[0], ("tiled", 85, None))
+            self.assertNotIn(("shrink_spatial", None), calls)
+        finally:
+            latent._upscale_video_model_tiled = original_tiled
+            latent._shrink_model_tile_param = original_shrink_model
+
    def test_interp_temporal_params_preserve_upstream_defaults(self):
        latent = importlib.import_module("dumas_h3_latent_upscale")
        chunk_length, temporal_overlap = latent._effective_temporal_params({