From 0f035dbb83d37b4edeffa16816f46c379151ef0e Mon Sep 17 00:00:00 2001 From: MSXYZ-GenAI Date: Sat, 25 Apr 2026 21:28:23 +0300 Subject: [PATCH] fix: add OOM fallback for tiled DLAA --- VideoTAADLAA.py | 41 ++++++++++++++++++++++++++++++++++------- 1 file changed, 34 insertions(+), 7 deletions(-) diff --git a/VideoTAADLAA.py b/VideoTAADLAA.py index ca9e7af..5353036 100644 --- a/VideoTAADLAA.py +++ b/VideoTAADLAA.py @@ -309,7 +309,7 @@ class VideoTAADLAA: if vram_mb <= 8192: tile_size = 512 elif vram_mb <= 16384: tile_size = 1024 - else: tile_size = 2048 + else: tile_size = 1024 # Estimate tile count if H > tile_size or W > tile_size: @@ -337,12 +337,39 @@ class VideoTAADLAA: # DLAA refinement pass if dlaa_strength > 0: - dlaa_out = self._tiled_forward( - net, - rgb, - tile_size=tile_size, - overlap=32 - ) + try: + dlaa_out = self._tiled_forward( + net, + rgb, + tile_size=tile_size, + overlap=32 + ) + except torch.OutOfMemoryError: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + logger.warning("[DLAA] OOM detected, retrying with smaller tiles") + + retry_tile = max(256, tile_size // 2) + + try: + dlaa_out = self._tiled_forward( + net, + rgb, + tile_size=retry_tile, + overlap=32 + ) + except torch.OutOfMemoryError: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + retry_tile = 256 + dlaa_out = self._tiled_forward( + net, + rgb, + tile_size=retry_tile, + overlap=32 + ) # Global luminance