From 179fa0778ea8a73fc9a20803f7f8be12ae49d76d Mon Sep 17 00:00:00 2001
From: Chris Dumas <chrisldumas@hotmail.com>
Date: Fri, 4 Sep 2026 07:33:10 +0000
Subject: [PATCH] Add CPU fallback for latent upscale OOM

---
 dumas_h3_latent_upscale.py        | 14 ++++++++++++
 tests/test_dumas_h3_longvideos.py | 36 +++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/dumas_h3_latent_upscale.py b/dumas_h3_latent_upscale.py
index f88d4d8..f51fa1b 100644
--- a/dumas_h3_latent_upscale.py
+++ b/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()
diff --git a/tests/test_dumas_h3_longvideos.py b/tests/test_dumas_h3_longvideos.py
index e804ba3..330837b 100644
--- a/tests/test_dumas_h3_longvideos.py
+++ b/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"