Harden latent upscale GPU cleanup

This commit is contained in:
2026-09-04 10:05:24 +00:00
parent cbabcf8208
commit 618a48e4d9
2 changed files with 45 additions and 16 deletions
+10 -3
View File
@@ -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: