Add CPU fallback for latent upscale OOM

dumas_h3_latent_upscale.py

@@ -602,6 +602,20 @@ def upscale_video_model(video, param):
             raise
         smaller = _shrink_model_tile_param(param)
         if smaller is None:
+            if str(param.get("device", "cuda")) != "cpu":
+                cpu_param = dict(param)
+                cpu_param["device"] = "cpu"
+                cpu_param["precision"] = "fp32"
+                try:
+                    gc.collect()
+                except Exception:
+                    pass
+                if torch.cuda.is_available():
+                    try:
+                        torch.cuda.empty_cache()
+                    except Exception:
+                        pass
+                return upscale_video_model(video, cpu_param)
             raise
         try:
             gc.collect()

tests/test_dumas_h3_longvideos.py

@@ -1196,6 +1196,42 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
         self.assertEqual(smaller["chunk_length"], 17)
         self.assertEqual(smaller["temporal_overlap"], 0)
 
+    def test_upscale_video_model_falls_back_to_cpu_when_gpu_cannot_shrink(self):
+        latent = importlib.import_module("dumas_h3_latent_upscale")
+        calls = []
+
+        original_tiled = latent._upscale_video_model_tiled
+        original_shrink = latent._shrink_model_tile_param
+        original_is_available = latent.torch.cuda.is_available
+        original_empty_cache = latent.torch.cuda.empty_cache
+        original_gc_collect = latent.gc.collect
+        try:
+            latent._shrink_model_tile_param = lambda _param: None
+            latent.torch.cuda.is_available = lambda: True
+            latent.torch.cuda.empty_cache = lambda: None
+            latent.gc.collect = lambda: None
+
+            def fake_tiled(video, param):
+                calls.append(dict(param))
+                if param.get("device") == "cpu":
+                    return "ok", 8, 16
+                raise RuntimeError("out of memory")
+
+            latent._upscale_video_model_tiled = fake_tiled
+
+            result = latent.upscale_video_model("video", {"device": "cuda", "precision": "fp16"})
+
+            self.assertEqual(result, ("ok", 8, 16))
+            self.assertEqual(calls[0]["device"], "cuda")
+            self.assertEqual(calls[1]["device"], "cpu")
+            self.assertEqual(calls[1]["precision"], "fp32")
+        finally:
+            latent._upscale_video_model_tiled = original_tiled
+            latent._shrink_model_tile_param = original_shrink
+            latent.torch.cuda.is_available = original_is_available
+            latent.torch.cuda.empty_cache = original_empty_cache
+            latent.gc.collect = original_gc_collect
+
     def test_compose_persistent_does_not_expand_ambiguous_plural_to_full_cast(self):
         active = self.module.parse_wardrobe(
             "Maya = she, red jacket\n"