Retry temporal split after latent upscale OOM

dumas_h3_latent_upscale.py

@@ -64,12 +64,36 @@ def _effective_temporal_params(param, frame_count=None):
         if chunk_length <= 0:
             chunk_length = 85
             temporal_overlap = 17
-        if frame_count is not None and CUDA_MODEL_CHUNK_LENGTH < int(frame_count) <= chunk_length:
-            chunk_length = CUDA_MODEL_CHUNK_LENGTH
-            temporal_overlap = 0
     return chunk_length, temporal_overlap
 
 
+def _is_oom_error(exc):
+    text = str(exc).lower()
+    return "out of memory" in text or "exhausted its gpu spatial fallbacks" in text
+
+
+def _retry_with_smaller_temporal(video, param, upscaler, exc):
+    if not _is_oom_error(exc):
+        raise exc
+    smaller = _shrink_temporal_param(param)
+    if smaller is None:
+        raise exc
+    LOGGER.info(
+        "H3 latent upscale temporal OOM: retrying with chunk_length=%s overlap=%s",
+        smaller.get("chunk_length"), smaller.get("temporal_overlap"),
+    )
+    try:
+        gc.collect()
+    except Exception:
+        pass
+    if torch.cuda.is_available():
+        try:
+            torch.cuda.empty_cache()
+        except Exception:
+            pass
+    return _upscale_video_temporal_chunks(video, smaller, upscaler)
+
+
 def _models_dir():
     try:
         if _LATENT_UPSCALE_FOLDER not in folder_paths.folder_names_and_paths:
@@ -807,7 +831,10 @@ def _upscale_video_temporal_chunks(video, param, upscaler):
             "H3 latent upscale: single temporal chunk, tokens=%s frames=%s chunk_length=%s overlap=%s",
             t, frame_count, chunk_length, temporal_overlap,
         )
+        try:
             return upscaler(video, chunk_param)
+        except RuntimeError as exc:
+            return _retry_with_smaller_temporal(video, chunk_param, upscaler, exc)
 
     bounds = _temporal_segments(t, chunk_length, temporal_overlap)
     if len(bounds) <= 1:
@@ -815,7 +842,10 @@ def _upscale_video_temporal_chunks(video, param, upscaler):
             "H3 latent upscale: single temporal segment, tokens=%s frames=%s chunk_length=%s overlap=%s",
             t, frame_count, chunk_length, temporal_overlap,
         )
+        try:
             return upscaler(video, chunk_param)
+        except RuntimeError as exc:
+            return _retry_with_smaller_temporal(video, chunk_param, upscaler, exc)
 
     orig_dtype = video.dtype
     out = None
@@ -833,25 +863,7 @@ def _upscale_video_temporal_chunks(video, param, upscaler):
         try:
             chunk_out, chunk_h, chunk_w = upscaler(chunk, chunk_param)
         except RuntimeError as exc:
-            if "out of memory" not in str(exc).lower():
-                raise
-            smaller = _shrink_temporal_param(chunk_param)
-            if smaller is None:
-                raise
-            LOGGER.info(
-                "H3 latent upscale temporal OOM: retrying with chunk_length=%s overlap=%s",
-                smaller.get("chunk_length"), smaller.get("temporal_overlap"),
-            )
-            try:
-                gc.collect()
-            except Exception:
-                pass
-            if torch.cuda.is_available():
-                try:
-                    torch.cuda.empty_cache()
-                except Exception:
-                    pass
-            return _upscale_video_temporal_chunks(video, smaller, upscaler)
+            return _retry_with_smaller_temporal(video, chunk_param, upscaler, exc)
         chunk_out = chunk_out.to(device="cpu", dtype=orig_dtype)
         if out is None:
             out_h, out_w = chunk_h, chunk_w

tests/test_dumas_h3_longvideos.py

@@ -1273,7 +1273,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
         self.assertEqual(chunk_length, 85)
         self.assertEqual(temporal_overlap, 17)
 
-    def test_cuda_model_temporal_params_cap_single_chunk_saved_workflows(self):
+    def test_cuda_model_temporal_params_keep_short_saved_workflows_until_oom(self):
         latent = importlib.import_module("dumas_h3_latent_upscale")
         chunk_length, temporal_overlap = latent._effective_temporal_params({
             "mode": "model",
@@ -1281,8 +1281,8 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
             "chunk_length": 85,
             "temporal_overlap": 17,
         }, frame_count=85)
-        self.assertEqual(chunk_length, 17)
-        self.assertEqual(temporal_overlap, 0)
+        self.assertEqual(chunk_length, 85)
+        self.assertEqual(temporal_overlap, 17)
 
     def test_interp_temporal_params_preserve_upstream_defaults(self):
         latent = importlib.import_module("dumas_h3_latent_upscale")