Add temporal OOM backoff for latent upscale

dumas_h3_latent_upscale.py

@@ -502,6 +502,43 @@ def _shrink_model_tile_param(param):
     return next_param
 
 
+def _shrink_temporal_param(param):
+    def _lower_17(v):
+        v = int(v or 0)
+        if v <= 0:
+            return 0
+        v = (v // 2 // 17) * 17
+        return max(0, v)
+
+    next_param = dict(param)
+    chunk_length = int(next_param.get("chunk_length", 0) or 0)
+    temporal_overlap = int(next_param.get("temporal_overlap", 0) or 0)
+    if chunk_length <= 17:
+        return None
+
+    new_chunk_length = _lower_17(chunk_length)
+    if new_chunk_length < 17:
+        new_chunk_length = chunk_length - 17
+    if new_chunk_length < 17:
+        return None
+    if new_chunk_length >= chunk_length:
+        new_chunk_length = chunk_length - 17
+        if new_chunk_length < 17:
+            return None
+
+    new_overlap = min(_lower_17(temporal_overlap), max(0, new_chunk_length - 17))
+    if new_overlap >= new_chunk_length:
+        new_overlap = max(0, new_chunk_length - 17)
+    if new_overlap >= new_chunk_length:
+        new_overlap = 0
+    if new_overlap >= new_chunk_length:
+        return None
+
+    next_param["chunk_length"] = new_chunk_length
+    next_param["temporal_overlap"] = new_overlap
+    return next_param
+
+
 def _upscale_video_model_core(video, param):
     model_name = param["model_name"]
     device = param.get("device", "cuda")
@@ -687,7 +724,24 @@ def _upscale_video_temporal_chunks(video, param, upscaler):
     out_h = out_w = None
     for i, (t0, t1) in enumerate(bounds):
         chunk = video[:, :, t0:t1].contiguous()
-        chunk_out, chunk_h, chunk_w = upscaler(chunk, param)
+        try:
+            chunk_out, chunk_h, chunk_w = upscaler(chunk, param)
+        except RuntimeError as exc:
+            if "out of memory" not in str(exc).lower():
+                raise
+            smaller = _shrink_temporal_param(param)
+            if smaller is None:
+                raise
+            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)
         chunk_out = chunk_out.to(device="cpu", dtype=orig_dtype)
         if out is None:
             out_h, out_w = chunk_h, chunk_w