Retry long videos decode with tiling fallback
This commit is contained in:
@@ -6379,8 +6379,21 @@ class H3LongVideos:
|
|||||||
# the same reason the decoded frames are.
|
# the same reason the decoded frames are.
|
||||||
decode_video_start = time.perf_counter()
|
decode_video_start = time.perf_counter()
|
||||||
shot_latent = _copy_sample_latent(refined_out)
|
shot_latent = _copy_sample_latent(refined_out)
|
||||||
|
try:
|
||||||
video = _decode_video(vae, refined_out, tiled, free_first=model,
|
video = _decode_video(vae, refined_out, tiled, free_first=model,
|
||||||
tile_t=decode_tile_frames, tile_xy=decode_tile_size)
|
tile_t=decode_tile_frames, tile_xy=decode_tile_size)
|
||||||
|
except Exception as e:
|
||||||
|
# Decode is the biggest allocation in the run. If the straight path
|
||||||
|
# fails, retry once with tiled decode and without the aggressive unload
|
||||||
|
# so a marginal card can still finish the render.
|
||||||
|
if not _is_oom(e) and "decode" not in str(e).lower():
|
||||||
|
raise
|
||||||
|
mm.soft_empty_cache(True)
|
||||||
|
retry_tiled = True
|
||||||
|
retry_tile_t = decode_tile_frames or 16
|
||||||
|
retry_tile_xy = decode_tile_size or 256
|
||||||
|
video = _decode_video(vae, refined_out, retry_tiled, free_first=None,
|
||||||
|
tile_t=retry_tile_t, tile_xy=retry_tile_xy)
|
||||||
timing["decode_video"] += time.perf_counter() - decode_video_start
|
timing["decode_video"] += time.perf_counter() - decode_video_start
|
||||||
decode_audio_start = time.perf_counter()
|
decode_audio_start = time.perf_counter()
|
||||||
audio = _decode_audio(audio_vae, out)
|
audio = _decode_audio(audio_vae, out)
|
||||||
|
|||||||
@@ -226,6 +226,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
|
|||||||
original_cleanup = self.module._deep_cleanup
|
original_cleanup = self.module._deep_cleanup
|
||||||
original_nested = getattr(self.module.comfy.nested_tensor, "NestedTensor", None)
|
original_nested = getattr(self.module.comfy.nested_tensor, "NestedTensor", None)
|
||||||
try:
|
try:
|
||||||
|
sentinel_model = object()
|
||||||
self.module.comfy.nested_tensor.NestedTensor = FakeNestedTensor
|
self.module.comfy.nested_tensor.NestedTensor = FakeNestedTensor
|
||||||
|
|
||||||
def common_ksampler(*args, **kwargs):
|
def common_ksampler(*args, **kwargs):
|
||||||
@@ -243,7 +244,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
|
|||||||
self.module._deep_cleanup = lambda: None
|
self.module._deep_cleanup = lambda: None
|
||||||
|
|
||||||
result = self.module.H3LongVideos()._render(
|
result = self.module.H3LongVideos()._render(
|
||||||
model=object(),
|
model=sentinel_model,
|
||||||
clip=types.SimpleNamespace(
|
clip=types.SimpleNamespace(
|
||||||
tokenize=lambda text, **kwargs: text,
|
tokenize=lambda text, **kwargs: text,
|
||||||
encode_from_tokens_scheduled=lambda tokens: tokens,
|
encode_from_tokens_scheduled=lambda tokens: tokens,
|
||||||
@@ -289,6 +290,95 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
|
|||||||
else:
|
else:
|
||||||
self.module.comfy.nested_tensor.NestedTensor = original_nested
|
self.module.comfy.nested_tensor.NestedTensor = original_nested
|
||||||
|
|
||||||
|
def test_render_retries_decode_with_tiling_after_decode_oom(self):
|
||||||
|
class FakeTensor:
|
||||||
|
def __init__(self, name):
|
||||||
|
self.name = name
|
||||||
|
|
||||||
|
def detach(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
def to(self, *args, **kwargs):
|
||||||
|
return self
|
||||||
|
|
||||||
|
class FakeNestedTensor:
|
||||||
|
def __init__(self, parts):
|
||||||
|
self._parts = tuple(parts)
|
||||||
|
self.is_nested = True
|
||||||
|
|
||||||
|
def unbind(self):
|
||||||
|
return self._parts
|
||||||
|
|
||||||
|
decode_calls = []
|
||||||
|
cleanup_calls = []
|
||||||
|
first_out = {"samples": FakeNestedTensor((FakeTensor("v1"), FakeTensor("a1")))}
|
||||||
|
|
||||||
|
original_common_ksampler = self.module.nodes.common_ksampler
|
||||||
|
original_build = self.module._build_shot_conditioning
|
||||||
|
original_evict = self.module._evict_all_but
|
||||||
|
original_decode_video = self.module._decode_video
|
||||||
|
original_decode_audio = self.module._decode_audio
|
||||||
|
original_cleanup = self.module._deep_cleanup
|
||||||
|
original_soft_empty_cache = self.module.mm.soft_empty_cache
|
||||||
|
original_nested = getattr(self.module.comfy.nested_tensor, "NestedTensor", None)
|
||||||
|
try:
|
||||||
|
sentinel_model = object()
|
||||||
|
self.module.comfy.nested_tensor.NestedTensor = FakeNestedTensor
|
||||||
|
self.module.nodes.common_ksampler = lambda *args, **kwargs: (first_out,)
|
||||||
|
self.module._build_shot_conditioning = lambda *_args, **_kwargs: (
|
||||||
|
"cond",
|
||||||
|
{"samples": FakeNestedTensor((FakeTensor("basev"), FakeTensor("basea")))},
|
||||||
|
)
|
||||||
|
self.module._evict_all_but = lambda *_args, **_kwargs: None
|
||||||
|
|
||||||
|
def decode_video(_vae, out_latent, tiled, free_first=None, tile_t=None, tile_xy=None):
|
||||||
|
decode_calls.append((tiled, free_first, tile_t, tile_xy))
|
||||||
|
if len(decode_calls) == 1:
|
||||||
|
raise RuntimeError("CUDA out of memory during decode")
|
||||||
|
return out_latent
|
||||||
|
|
||||||
|
self.module._decode_video = decode_video
|
||||||
|
self.module._decode_audio = lambda _vae, out_latent: out_latent
|
||||||
|
self.module.mm.soft_empty_cache = lambda *args, **kwargs: cleanup_calls.append((args, kwargs))
|
||||||
|
self.module._deep_cleanup = lambda: None
|
||||||
|
|
||||||
|
result = self.module.H3LongVideos()._render(
|
||||||
|
model=sentinel_model,
|
||||||
|
clip=types.SimpleNamespace(
|
||||||
|
tokenize=lambda text, **kwargs: text,
|
||||||
|
encode_from_tokens_scheduled=lambda tokens: tokens,
|
||||||
|
),
|
||||||
|
vae=object(),
|
||||||
|
audio_vae=object(),
|
||||||
|
negative="negative",
|
||||||
|
prompt="beat",
|
||||||
|
w=128,
|
||||||
|
h=64,
|
||||||
|
ln=24,
|
||||||
|
fps=24,
|
||||||
|
tiled=False,
|
||||||
|
sa=(123, 20, 1.0, "res_multistep", "simple", 1.0),
|
||||||
|
handoff=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(len(decode_calls), 2)
|
||||||
|
self.assertEqual(decode_calls[0], (False, sentinel_model, 0, 0))
|
||||||
|
self.assertEqual(decode_calls[1], (True, None, 16, 256))
|
||||||
|
self.assertTrue(cleanup_calls)
|
||||||
|
self.assertIs(result[1], first_out)
|
||||||
|
finally:
|
||||||
|
self.module.nodes.common_ksampler = original_common_ksampler
|
||||||
|
self.module._build_shot_conditioning = original_build
|
||||||
|
self.module._evict_all_but = original_evict
|
||||||
|
self.module._decode_video = original_decode_video
|
||||||
|
self.module._decode_audio = original_decode_audio
|
||||||
|
self.module._deep_cleanup = original_cleanup
|
||||||
|
self.module.mm.soft_empty_cache = original_soft_empty_cache
|
||||||
|
if original_nested is None:
|
||||||
|
delattr(self.module.comfy.nested_tensor, "NestedTensor")
|
||||||
|
else:
|
||||||
|
self.module.comfy.nested_tensor.NestedTensor = original_nested
|
||||||
|
|
||||||
def test_detail_pass_treats_falsey_strings_as_disabled(self):
|
def test_detail_pass_treats_falsey_strings_as_disabled(self):
|
||||||
calls = []
|
calls = []
|
||||||
original_common_ksampler = self.module.nodes.common_ksampler
|
original_common_ksampler = self.module.nodes.common_ksampler
|
||||||
|
|||||||
Reference in New Issue
Block a user