Back off latent upscale tile size on OOM

dumas_h3_latent_upscale.py

@@ -419,6 +419,36 @@ def _compute_spatial_grid(h, w, th, tw, ol_h, ol_w, min_th=0, min_tw=0):
     return rows, cols, trows, tcols, row_ovl, col_ovl
 
 
+def _shrink_model_tile_param(param):
+    next_param = dict(param)
+    mode = str(next_param.get("tile_size_mode") or "specific_size")
+    if mode == "rows_cols":
+        rows = min(9, max(1, int(next_param.get("grid_rows", 2) or 2) + 1))
+        cols = min(9, max(1, int(next_param.get("grid_cols", 2) or 2) + 1))
+        if rows == int(next_param.get("grid_rows", 2) or 2) and cols == int(next_param.get("grid_cols", 2) or 2):
+            return None
+        next_param["grid_rows"] = rows
+        next_param["grid_cols"] = cols
+        next_param["spatial_w_overlap"] = max(0, int(next_param.get("spatial_w_overlap", 0) or 0) // 2)
+        next_param["spatial_h_overlap"] = max(0, int(next_param.get("spatial_h_overlap", 0) or 0) // 2)
+        next_param["fade_width"] = max(0, int(next_param.get("fade_width", 0) or 0) // 2)
+        next_param["fade_height"] = max(0, int(next_param.get("fade_height", 0) or 0) // 2)
+        return next_param
+
+    tile_w = int(next_param.get("tile_width", 0) or 0)
+    tile_h = int(next_param.get("tile_height", 0) or 0)
+    new_w = max(32, (max(32, tile_w) // 2 // 32) * 32)
+    new_h = max(32, (max(32, tile_h) // 2 // 32) * 32)
+    if new_w >= tile_w and new_h >= tile_h:
+        return None
+    next_param["tile_width"] = new_w
+    next_param["tile_height"] = new_h
+    next_param["overlap"] = max(0, int(next_param.get("overlap", 0) or 0) // 2)
+    next_param["fade_width"] = max(0, int(next_param.get("fade_width", 0) or 0) // 2)
+    next_param["fade_height"] = max(0, int(next_param.get("fade_height", 0) or 0) // 2)
+    return next_param
+
+
 def _upscale_video_model_core(video, param):
     model_name = param["model_name"]
     device = param.get("device", "cuda")
@@ -456,6 +486,18 @@ def _upscale_video_model_core(video, param):
 
 
 def upscale_video_model(video, param):
+    try:
+        return _upscale_video_model_tiled(video, param)
+    except RuntimeError as exc:
+        if "out of memory" not in str(exc).lower():
+            raise
+        smaller = _shrink_model_tile_param(param)
+        if smaller is None:
+            raise
+        return upscale_video_model(video, smaller)
+
+
+def _upscale_video_model_tiled(video, param):
     _, _, _, h_in, w_in = video.shape
     h_out, w_out = _resolve_target_size(param, h_in, w_in)
     mode = str(param.get("tile_size_mode") or "specific_size")

tests/test_dumas_h3_longvideos.py

@@ -1160,6 +1160,20 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
         self.assertIs(tagged, exc)
         self.assertEqual(getattr(tagged, "_h3_stage", ""), "latent_upscale")
 
+    def test_shrink_model_tile_param_reduces_tile_size(self):
+        latent = importlib.import_module("dumas_h3_latent_upscale")
+        smaller = latent._shrink_model_tile_param({
+            "tile_size_mode": "specific_size",
+            "tile_width": 512,
+            "tile_height": 512,
+            "overlap": 64,
+            "fade_width": 32,
+            "fade_height": 32,
+        })
+        self.assertIsNotNone(smaller)
+        self.assertLess(smaller["tile_width"], 512)
+        self.assertLess(smaller["tile_height"], 512)
+
     def test_compose_persistent_does_not_expand_ambiguous_plural_to_full_cast(self):
         active = self.module.parse_wardrobe(
             "Maya = she, red jacket\n"