From 618a48e4d97050ed580aacc98177bdccc18abe64 Mon Sep 17 00:00:00 2001
From: Chris Dumas <chrisldumas@hotmail.com>
Date: Fri, 4 Sep 2026 10:05:24 +0000
Subject: [PATCH] Harden latent upscale GPU cleanup

---
 dumas_h3_longvideos.py            | 48 ++++++++++++++++++++++---------
 tests/test_dumas_h3_longvideos.py | 13 +++++++--
 2 files changed, 45 insertions(+), 16 deletions(-)

diff --git a/dumas_h3_longvideos.py b/dumas_h3_longvideos.py
index 28e88a6..1aad51e 100644
--- a/dumas_h3_longvideos.py
+++ b/dumas_h3_longvideos.py
@@ -6055,7 +6055,7 @@ def quant_accel_note(model):
 
 
 
-def _evict_all_but(keep_model):
+def _evict_all_but(keep_model):
     """Unload every model EXCEPT the diffusion model from the GPU.
 
     This is the fix for VRAM ratcheting across a long chain. soft_empty_cache()
@@ -6081,15 +6081,38 @@ def _evict_all_but(keep_model):
     except Exception:
         try:
             mm.soft_empty_cache(True)
-        except Exception:
-            pass
-
-
-
-
-
-
-class H3LongVideos:
+        except Exception:
+            pass
+
+
+def _evict_for_latent_upscale(model):
+    """Clear the sampler model before loading the auxiliary latent upscaler."""
+    try:
+        unload_clones = getattr(mm, "unload_model_and_clones", None)
+        if callable(unload_clones):
+            try:
+                unload_clones(model, unload_additional_models=False)
+                mm.soft_empty_cache()
+                return
+            except Exception:
+                pass
+        mm.unload_all_models()
+    except Exception:
+        pass
+    try:
+        mm.soft_empty_cache(True)
+    except Exception:
+        try:
+            mm.soft_empty_cache()
+        except Exception:
+            pass
+
+
+
+
+
+
+class H3LongVideos:
     CATEGORY = "Dumas/MiniMax"
     FUNCTION = "run"
     # fps is emitted as BOTH types on purpose: ComfyUI does not coerce between them,
@@ -6662,9 +6685,8 @@ class H3LongVideos:
                 parts = _nested_tensor_parts(out_samples)
                 if not getattr(out_samples, "is_nested", False) or len(parts) < 2:
                     raise RuntimeError("latent upscale expects a nested AV latent")
-                if latent_upscale_mode == "model" and str(latent_upscale_param.get("device", "cuda")) == "cuda" and hasattr(model, "clone_base_uuid"):
-                    mm.unload_model_and_clones(model, unload_additional_models=False)
-                    mm.soft_empty_cache()
+                if latent_upscale_mode == "model" and str(latent_upscale_param.get("device", "cuda")) == "cuda":
+                    _evict_for_latent_upscale(model)
                 upscaled_video, up_h, up_w = _upscale_latent_video(parts[0], latent_upscale_param)
                 full_audio = parts[1]
                 # Drop the first-pass sampling state before we start the refinement
diff --git a/tests/test_dumas_h3_longvideos.py b/tests/test_dumas_h3_longvideos.py
index 4539299..55249a1 100644
--- a/tests/test_dumas_h3_longvideos.py
+++ b/tests/test_dumas_h3_longvideos.py
@@ -374,6 +374,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
         original_upscale = self.module._upscale_latent_video
         original_copy_sample = self.module._copy_sample_latent
         original_unload = getattr(self.module.mm, "unload_model_and_clones", None)
+        original_unload_all = self.module.mm.unload_all_models
         original_nested = getattr(self.module.comfy.nested_tensor, "NestedTensor", None)
         try:
             self.module.comfy.nested_tensor.NestedTensor = FakeNestedTensor
@@ -397,7 +398,11 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
                 order.append("cleanup")
 
             def unload_model_and_clones(*_args, **_kwargs):
-                order.append("unload_h3")
+                order.append("unload_h3_failed")
+                raise RuntimeError("model wrapper does not expose clone metadata")
+
+            def unload_all_models(*_args, **_kwargs):
+                order.append("unload_all")
 
             def upscale_latent_video(video, param):
                 order.append("upscale")
@@ -410,6 +415,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
             )
             self.module._evict_all_but = lambda *_args, **_kwargs: None
             self.module.mm.unload_model_and_clones = unload_model_and_clones
+            self.module.mm.unload_all_models = unload_all_models
             self.module._upscale_latent_video = upscale_latent_video
             self.module._copy_sample_latent = lambda sampled: sampled["samples"].unbind()
             self.module._decode_video = decode_video
@@ -417,7 +423,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
             self.module._deep_cleanup = cleanup
 
             self.module.H3LongVideos()._render(
-                model=types.SimpleNamespace(clone_base_uuid="h3"),
+                model=object(),
                 clip=types.SimpleNamespace(
                     tokenize=lambda text, **kwargs: text,
                     encode_from_tokens_scheduled=lambda tokens: tokens,
@@ -447,7 +453,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
             )
 
             self.assertEqual(order[0], "sample")
-            self.assertLess(order.index("unload_h3"), order.index("upscale"))
+            self.assertLess(order.index("unload_all"), order.index("upscale"))
             self.assertLess(order.index("upscale"), order.index("latent_upscale_sample"))
             self.assertLess(order.index("audio"), order.index("video"))
             self.assertEqual(order[-1], "cleanup")
@@ -464,6 +470,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
                 delattr(self.module.mm, "unload_model_and_clones")
             else:
                 self.module.mm.unload_model_and_clones = original_unload
+            self.module.mm.unload_all_models = original_unload_all
             if original_nested is None:
                 delattr(self.module.comfy.nested_tensor, "NestedTensor")
             else: