Harden latent upscale GPU cleanup
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user