797 lines
31 KiB
Python
797 lines
31 KiB
Python
"""
|
|
Universal Smart VAE Decode - v11.2 Production Hardened
|
|
Grok v11.0 + Claude refinements + Copilot security fixes
|
|
|
|
CRITICAL FIXES in v11.2:
|
|
- ZeroDivisionError protection in time scale detection
|
|
- OOM retry limit (prevents infinite loops)
|
|
- Thread-safe VAE cache
|
|
- Proper logging instead of print
|
|
- Atomic temp file handling
|
|
- Better error messages with recovery hints
|
|
|
|
Credits: Grok (disk offload, orientation), Kimi (memory safety), Claude (stability),
|
|
Gemini (math), GPT (structure), Copilot (security review)
|
|
Version: 11.2.0
|
|
License: MIT
|
|
GitHub: https://github.com/uczensokratesa/ComfyUI-Grok-SmartVAE
|
|
"""
|
|
|
|
import torch
|
|
import comfy.utils
|
|
from comfy.model_management import throw_exception_if_processing_interrupted
|
|
import gc
|
|
import math
|
|
import os
|
|
import tempfile
|
|
import warnings
|
|
import atexit
|
|
import logging
|
|
from typing import Optional, Tuple
|
|
|
|
# Setup logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
import psutil
|
|
PSUTIL_AVAILABLE = True
|
|
except ImportError:
|
|
PSUTIL_AVAILABLE = False
|
|
warnings.warn(
|
|
"psutil not available. Install with 'pip install psutil' for automatic RAM detection.\n"
|
|
"Disk offload will still work based on max_ram_frames threshold."
|
|
)
|
|
|
|
|
|
class UniversalSmartVAEDecode:
|
|
"""
|
|
Production-grade VAE decoder with disk offloading for massive workflows.
|
|
|
|
Features:
|
|
- Handles 2000+ frames with automatic memory management
|
|
- Disk offload when RAM insufficient
|
|
- Frame-perfect audio sync
|
|
- Multi-stage OOM recovery
|
|
- Thread-safe caching
|
|
"""
|
|
|
|
# Class-level constants
|
|
MAX_OOM_RETRIES = 5
|
|
DISK_MERGE_CHUNK_SIZE = 10
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
"frames_per_batch": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 128,
|
|
"step": 1,
|
|
"tooltip": "Frames to decode per batch. Auto-reduces on OOM."
|
|
}),
|
|
},
|
|
"optional": {
|
|
"overlap_frames": ("INT", {
|
|
"default": 2,
|
|
"min": 0,
|
|
"max": 16,
|
|
"step": 1,
|
|
"tooltip": "Temporal overlap for seamless stitching."
|
|
}),
|
|
"force_time_scale": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 16,
|
|
"step": 1,
|
|
"tooltip": "Manual time scale override (0=auto). E.g., 8 for LTX-Video."
|
|
}),
|
|
"enable_tiling": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Force spatial tiling. Auto-enables on OOM."
|
|
}),
|
|
"tile_size": ("INT", {
|
|
"default": 512,
|
|
"min": 256,
|
|
"max": 2048,
|
|
"step": 64,
|
|
"tooltip": "Tile size in pixels for spatial tiling."
|
|
}),
|
|
"verbose": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Show detailed progress logs."
|
|
}),
|
|
"max_ram_frames": ("INT", {
|
|
"default": 500,
|
|
"min": 100,
|
|
"max": 10000,
|
|
"step": 100,
|
|
"tooltip": "Threshold for disk offload. Tune to system RAM: 16GB=300, 32GB=600, 64GB=1000."
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "decode"
|
|
CATEGORY = "latent/video"
|
|
|
|
def __init__(self):
|
|
# Thread-safe per-VAE cache (dict instead of single values)
|
|
self._time_scale_cache = {}
|
|
self._force_scale_cache = {}
|
|
self.temp_dir = None
|
|
|
|
# Register cleanup
|
|
atexit.register(self._cleanup_temp_dir)
|
|
|
|
def _cleanup_temp_dir(self):
|
|
"""Cleanup temporary files (called on exit or error)."""
|
|
if self.temp_dir and os.path.exists(self.temp_dir):
|
|
try:
|
|
for file in os.listdir(self.temp_dir):
|
|
try:
|
|
os.remove(os.path.join(self.temp_dir, file))
|
|
except Exception as e:
|
|
logger.debug(f"Failed to remove temp file: {e}")
|
|
os.rmdir(self.temp_dir)
|
|
logger.debug(f"Cleaned up temp dir: {self.temp_dir}")
|
|
except Exception as e:
|
|
logger.debug(f"Temp dir cleanup error: {e}")
|
|
|
|
def __del__(self):
|
|
"""Backup cleanup on object destruction."""
|
|
self._cleanup_temp_dir()
|
|
|
|
def _get_available_vram(self) -> Optional[float]:
|
|
"""Get available VRAM in GB. Returns None on CPU or error."""
|
|
try:
|
|
if not torch.cuda.is_available():
|
|
return None
|
|
free_vram, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
|
|
return free_vram / (1024 ** 3)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to get VRAM info: {e}")
|
|
return None
|
|
|
|
def _get_available_ram(self) -> Optional[float]:
|
|
"""Get available system RAM in GB. Returns None if psutil unavailable."""
|
|
if PSUTIL_AVAILABLE:
|
|
try:
|
|
return psutil.virtual_memory().available / (1024 ** 3)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to get RAM info: {e}")
|
|
return None
|
|
|
|
def detect_time_scale(self, vae, latents: torch.Tensor, force_scale: int = 0,
|
|
verbose: bool = True) -> int:
|
|
"""
|
|
Detect temporal upsampling scale with ZeroDivisionError protection.
|
|
|
|
Priority: force > cache > metadata > empirical test > fallback
|
|
|
|
Args:
|
|
vae: VAE model
|
|
latents: Input latent tensor [B, C, F, H, W]
|
|
force_scale: User override (0 = auto)
|
|
verbose: Log detection process
|
|
|
|
Returns:
|
|
Detected time scale (integer >= 1)
|
|
"""
|
|
vae_id = id(vae)
|
|
|
|
# Priority 1: User override
|
|
if force_scale > 0:
|
|
if self._force_scale_cache.get(vae_id) != force_scale:
|
|
# Invalidate time scale cache if force changed
|
|
self._time_scale_cache.pop(vae_id, None)
|
|
self._force_scale_cache[vae_id] = force_scale
|
|
|
|
if verbose:
|
|
logger.info(f"🔧 Forced time scale: {force_scale}x")
|
|
|
|
self._time_scale_cache[vae_id] = force_scale
|
|
return force_scale
|
|
|
|
# Clear force scale tracking
|
|
self._force_scale_cache.pop(vae_id, None)
|
|
|
|
# Priority 2: Cache check
|
|
if vae_id in self._time_scale_cache:
|
|
return self._time_scale_cache[vae_id]
|
|
|
|
# Priority 3: VAE metadata
|
|
if hasattr(vae, 'downscale_index_formula') and vae.downscale_index_formula:
|
|
try:
|
|
time_scale = int(vae.downscale_index_formula[0])
|
|
if verbose:
|
|
logger.info(f"🔍 VAE metadata time scale: {time_scale}x")
|
|
self._time_scale_cache[vae_id] = time_scale
|
|
return time_scale
|
|
except Exception as e:
|
|
logger.debug(f"Metadata parsing failed: {e}")
|
|
|
|
# Priority 4: Empirical test (with ZeroDivisionError protection)
|
|
try:
|
|
total_latent_frames = latents.shape[2]
|
|
test_frames = min(5, total_latent_frames)
|
|
|
|
# CRITICAL FIX: Prevent division by zero
|
|
if test_frames <= 1:
|
|
if verbose:
|
|
logger.warning(
|
|
f"Not enough frames ({test_frames}) for empirical detection. "
|
|
f"Using fallback: 1x"
|
|
)
|
|
self._time_scale_cache[vae_id] = 1
|
|
return 1
|
|
|
|
test_sample = latents[:, :, 0:test_frames, :16, :16]
|
|
|
|
with torch.no_grad():
|
|
test_output = vae.decode(test_sample)
|
|
|
|
test_output = self._normalize_output(test_output, aspect_ratio=1.0)
|
|
output_frames = test_output.shape[0]
|
|
|
|
# Formula: output = 1 + (input - 1) * scale
|
|
# Solving: scale = (output - 1) / (input - 1)
|
|
# SAFE: test_frames > 1 guaranteed by check above
|
|
time_scale = max(1, (output_frames - 1) // (test_frames - 1))
|
|
|
|
if verbose:
|
|
logger.info(f"🔍 Auto-detected time scale: {time_scale}x ({test_frames}→{output_frames} frames)")
|
|
|
|
self._time_scale_cache[vae_id] = time_scale
|
|
|
|
# Cleanup
|
|
del test_output, test_sample
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
return time_scale
|
|
|
|
except Exception as e:
|
|
if verbose:
|
|
logger.warning(f"Time scale detection failed: {e}")
|
|
logger.info("Using safe fallback: 1x")
|
|
|
|
self._time_scale_cache[vae_id] = 1
|
|
return 1
|
|
|
|
def detect_output_size(self, vae, latents: torch.Tensor, h_latent: int, w_latent: int,
|
|
tile_size: int = 512, verbose: bool = True) -> Tuple[int, int]:
|
|
"""
|
|
Detect exact output H/W from single frame decode.
|
|
Falls back to tiled decode if standard OOMs.
|
|
|
|
Returns:
|
|
(output_height, output_width) tuple
|
|
"""
|
|
aspect_ratio = h_latent / float(w_latent)
|
|
test_sample = latents[:, :, 0:1, :, :]
|
|
|
|
try:
|
|
with torch.no_grad():
|
|
test_out = vae.decode(test_sample)
|
|
if verbose:
|
|
logger.info("🔍 Output size detected (standard decode)")
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e).lower():
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
try:
|
|
with torch.no_grad():
|
|
test_out = vae.decode_tiled(test_sample, tile_x=tile_size, tile_y=tile_size)
|
|
if verbose:
|
|
logger.info("🔍 Output size detected (tiled - OOM fallback)")
|
|
except Exception as inner_e:
|
|
raise RuntimeError("Failed to detect output size even with tiling.") from inner_e
|
|
else:
|
|
raise
|
|
|
|
test_out = self._normalize_output(test_out, aspect_ratio)
|
|
output_h, output_w = test_out.shape[1:3]
|
|
|
|
# Cleanup
|
|
del test_out, test_sample
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return output_h, output_w
|
|
|
|
def _estimate_chunk_vram(self, frames: int, channels: int, h: int, w: int,
|
|
time_scale: int = 1) -> float:
|
|
"""Estimate VRAM needed for chunk decode (GB)."""
|
|
spatial_scale = 8
|
|
latent_bytes = frames * channels * h * w * 4
|
|
output_frames = frames * time_scale
|
|
output_bytes = output_frames * 3 * (h * spatial_scale) * (w * spatial_scale) * 4
|
|
total_bytes = (latent_bytes + output_bytes) * 3.5 * 1.1
|
|
return total_bytes / (1024 ** 3)
|
|
|
|
def _estimate_output_ram(self, expected_frames: int, output_h: int, output_w: int) -> float:
|
|
"""Estimate total RAM for output tensor (GB)."""
|
|
bytes_per_frame = output_h * output_w * 3 * 4
|
|
return (expected_frames * bytes_per_frame) / (1024 ** 3)
|
|
|
|
def _normalize_output(self, tensor, aspect_ratio: Optional[float] = None) -> torch.Tensor:
|
|
"""
|
|
Normalize VAE output to [Frames, Height, Width, Channels].
|
|
|
|
Args:
|
|
tensor: VAE output (4D or 5D)
|
|
aspect_ratio: Latent h/w ratio for orientation detection
|
|
|
|
Returns:
|
|
Normalized [F, H, W, C] tensor
|
|
"""
|
|
# Handle list/tuple
|
|
if isinstance(tensor, (list, tuple)):
|
|
if not tensor:
|
|
raise ValueError("VAE returned empty output")
|
|
tensor = tensor[0]
|
|
|
|
if not isinstance(tensor, torch.Tensor):
|
|
raise TypeError(f"Expected Tensor, got {type(tensor)}")
|
|
|
|
# Ensure float32
|
|
if tensor.dtype != torch.float32:
|
|
tensor = tensor.float()
|
|
|
|
dim = tensor.dim()
|
|
|
|
if dim == 4:
|
|
# [N, C, H, W] or [N, H, W, C]
|
|
if tensor.shape[1] in (3, 4):
|
|
tensor = tensor.permute(0, 2, 3, 1)
|
|
|
|
elif dim == 5:
|
|
shape = list(tensor.shape)
|
|
|
|
# Squeeze B=1
|
|
if shape[0] == 1:
|
|
tensor = tensor.squeeze(0)
|
|
shape = list(tensor.shape)
|
|
|
|
# Find channel dim
|
|
try:
|
|
c_idx = next(i for i, s in enumerate(shape) if s in (3, 4))
|
|
except StopIteration:
|
|
raise ValueError(f"Cannot find channel dim (3 or 4) in shape {shape}")
|
|
|
|
# Remaining are [F, H, W]
|
|
remaining_idxs = [i for i in range(4) if i != c_idx]
|
|
remaining_sizes = [shape[i] for i in remaining_idxs]
|
|
|
|
# Sort: smallest=F, largest two=spatial
|
|
sorted_remaining = sorted(zip(remaining_idxs, remaining_sizes), key=lambda x: x[1])
|
|
f_idx = sorted_remaining[0][0]
|
|
spatial_large_idx = sorted_remaining[2][0]
|
|
spatial_small_idx = sorted_remaining[1][0]
|
|
|
|
# Assign H/W based on aspect ratio
|
|
if aspect_ratio is not None:
|
|
if aspect_ratio > 1.0:
|
|
# Portrait: H > W
|
|
h_idx = spatial_large_idx
|
|
w_idx = spatial_small_idx
|
|
else:
|
|
# Landscape: W > H
|
|
w_idx = spatial_large_idx
|
|
h_idx = spatial_small_idx
|
|
else:
|
|
# Default: portrait
|
|
h_idx = spatial_large_idx
|
|
w_idx = spatial_small_idx
|
|
|
|
# Permute to [F, H, W, C]
|
|
perm = [f_idx, h_idx, w_idx, c_idx]
|
|
tensor = tensor.permute(*perm)
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported dimension: {dim}D, shape {tensor.shape}")
|
|
|
|
# Clamp and convert RGBA→RGB
|
|
tensor = torch.clamp(tensor, 0.0, 1.0)
|
|
if tensor.shape[-1] > 3:
|
|
tensor = tensor[..., :3]
|
|
|
|
return tensor.contiguous()
|
|
|
|
def _center_crop_to_reference(self, tensor: torch.Tensor, h_ref: int, w_ref: int) -> torch.Tensor:
|
|
"""Center crop to reference dimensions."""
|
|
_, h, w, _ = tensor.shape
|
|
if h == h_ref and w == w_ref:
|
|
return tensor
|
|
|
|
h_offset = max(0, (h - h_ref) // 2)
|
|
w_offset = max(0, (w - w_ref) // 2)
|
|
|
|
return tensor[:, h_offset:h_offset + h_ref, w_offset:w_offset + w_ref, :]
|
|
|
|
def decode(self, vae, samples, frames_per_batch, overlap_frames=2, force_time_scale=0,
|
|
enable_tiling=False, tile_size=512, verbose=True, max_ram_frames=500):
|
|
"""
|
|
Main decode with disk offloading and OOM protection.
|
|
|
|
Features:
|
|
- Automatic disk offload for 700+ frame workflows
|
|
- Multi-stage OOM recovery with retry limits
|
|
- Frame-perfect audio sync
|
|
- Thread-safe operation
|
|
"""
|
|
latents = samples["samples"]
|
|
|
|
# ======== IMAGE PATH ========
|
|
if latents.dim() == 4:
|
|
if verbose:
|
|
logger.info("🖼️ Image decode")
|
|
|
|
with torch.no_grad():
|
|
if enable_tiling and hasattr(vae, 'decode_tiled'):
|
|
output = vae.decode_tiled(latents, tile_x=tile_size, tile_y=tile_size)
|
|
else:
|
|
output = vae.decode(latents)
|
|
|
|
return (self._normalize_output(output, aspect_ratio=1.0),)
|
|
|
|
# ======== VIDEO PATH ========
|
|
batch, channels, total_frames, h_latent, w_latent = latents.shape
|
|
|
|
if total_frames <= 0:
|
|
raise ValueError("Latent has no frames")
|
|
|
|
# Detect scales
|
|
time_scale = self.detect_time_scale(vae, latents, force_time_scale, verbose)
|
|
expected_frames = 1 + (total_frames - 1) * time_scale
|
|
aspect_ratio = h_latent / float(w_latent)
|
|
|
|
# Detect output size
|
|
output_h, output_w = self.detect_output_size(vae, latents, h_latent, w_latent, tile_size, verbose)
|
|
|
|
# RAM estimation
|
|
est_ram = self._estimate_output_ram(expected_frames, output_h, output_w)
|
|
available_ram = self._get_available_ram()
|
|
|
|
# Decide memory mode
|
|
use_disk_offload = (
|
|
expected_frames > max_ram_frames or
|
|
(available_ram is not None and est_ram > available_ram * 0.5)
|
|
)
|
|
|
|
# Setup temp dir
|
|
if use_disk_offload:
|
|
self.temp_dir = tempfile.mkdtemp(prefix="comfy_smartvae_")
|
|
if verbose:
|
|
logger.info(f"💾 Disk offload enabled ({expected_frames} frames, est {est_ram:.2f}GB RAM)")
|
|
|
|
# Display info
|
|
if verbose:
|
|
logger.info(f"🎬 Video decode:")
|
|
logger.info(f" Input: {total_frames} latent frames")
|
|
logger.info(f" Time scale: {time_scale}x")
|
|
logger.info(f" Expected output: ~{expected_frames} frames")
|
|
logger.info(f" Output resolution: {output_h}x{output_w}")
|
|
logger.info(f" Estimated RAM: {est_ram:.2f}GB")
|
|
|
|
if available_ram:
|
|
ram_pct = (est_ram / available_ram) * 100
|
|
logger.info(f" RAM usage: {ram_pct:.1f}% of {available_ram:.1f}GB available")
|
|
|
|
orientation = 'portrait' if aspect_ratio < 1.0 else 'landscape'
|
|
logger.info(f" Aspect ratio: {aspect_ratio:.3f} ({orientation})")
|
|
logger.info(f" Mode: {'💾 Disk offload' if use_disk_offload else '🚀 In-memory (pinned)'}")
|
|
|
|
if est_ram > 16:
|
|
warnings.warn(
|
|
f"High RAM usage ({est_ram:.1f}GB). "
|
|
f"Consider lowering resolution or increasing max_ram_frames."
|
|
)
|
|
|
|
# Validate parameters
|
|
frames_per_batch = max(1, min(frames_per_batch, total_frames))
|
|
overlap_frames = max(0, min(overlap_frames, frames_per_batch - 1))
|
|
initial_overlap = overlap_frames
|
|
|
|
# VRAM optimization
|
|
available_vram = self._get_available_vram()
|
|
if available_vram is not None:
|
|
chunk_frames = frames_per_batch + 2 * overlap_frames
|
|
est_vram = self._estimate_chunk_vram(chunk_frames, channels, h_latent, w_latent, time_scale)
|
|
|
|
if est_vram > available_vram * 0.55:
|
|
reduction = (available_vram * 0.45) / est_vram
|
|
old_batch = frames_per_batch
|
|
frames_per_batch = max(1, int(frames_per_batch * reduction))
|
|
overlap_frames = min(overlap_frames, frames_per_batch - 1)
|
|
|
|
if verbose:
|
|
logger.info(f"📉 Predictive VRAM reduction:")
|
|
logger.info(f" Batch: {old_batch} → {frames_per_batch}")
|
|
logger.info(f" Est: {est_vram:.2f}GB / Avail: {available_vram:.2f}GB")
|
|
|
|
if verbose:
|
|
logger.info(f" Batch size: {frames_per_batch}")
|
|
logger.info(f" Overlap: {overlap_frames} frames")
|
|
|
|
# Initialize storage
|
|
temp_files = [] if use_disk_offload else None
|
|
final_output = None
|
|
current_write_idx = 0
|
|
|
|
# Pre-allocation
|
|
if not use_disk_offload:
|
|
try:
|
|
final_output = torch.empty(
|
|
expected_frames, output_h, output_w, 3,
|
|
dtype=torch.float32,
|
|
device='cpu',
|
|
pin_memory=torch.cuda.is_available()
|
|
)
|
|
if verbose:
|
|
logger.info(" ✓ Pre-allocated pinned memory")
|
|
except Exception as e:
|
|
if verbose:
|
|
logger.warning(f"Pre-allocation failed: {e}")
|
|
logger.info("→ Falling back to disk offload")
|
|
use_disk_offload = True
|
|
self.temp_dir = tempfile.mkdtemp(prefix="comfy_smartvae_")
|
|
temp_files = []
|
|
|
|
# Processing state
|
|
current_batch = frames_per_batch
|
|
start_idx = 0
|
|
frames_processed = 0
|
|
last_frames_processed = -1
|
|
stagnation_count = 0
|
|
MAX_STAGNATION = 3
|
|
cumulative_output_frames = 0
|
|
oom_retry_count = 0 # CRITICAL FIX: OOM retry limit
|
|
|
|
pbar = comfy.utils.ProgressBar(total_frames)
|
|
|
|
# Adaptive logging
|
|
est_total_chunks = math.ceil(total_frames / frames_per_batch)
|
|
log_every_n = 5 if est_total_chunks > 20 else 1
|
|
chunk_count = 0
|
|
|
|
# ======== PROCESSING LOOP WITH RETRY LIMIT ========
|
|
while start_idx < total_frames:
|
|
throw_exception_if_processing_interrupted()
|
|
|
|
# CRITICAL FIX: OOM retry limit
|
|
if oom_retry_count >= self.MAX_OOM_RETRIES:
|
|
raise RuntimeError(
|
|
f"Exceeded maximum OOM retries ({self.MAX_OOM_RETRIES}).\n"
|
|
f"Current settings: batch={current_batch}, tile={tile_size}, tiling={enable_tiling}\n"
|
|
f"This indicates insufficient VRAM for this workflow.\n\n"
|
|
f"Suggestions:\n"
|
|
f" 1. Reduce latent resolution\n"
|
|
f" 2. Close other GPU applications\n"
|
|
f" 3. Process video in smaller segments"
|
|
)
|
|
|
|
# Stagnation guard
|
|
if frames_processed == last_frames_processed:
|
|
stagnation_count += 1
|
|
if stagnation_count >= MAX_STAGNATION:
|
|
raise RuntimeError(
|
|
f"Processing stalled at frame {start_idx}/{total_frames}.\n"
|
|
f"Settings: batch={current_batch}, overlap={overlap_frames}, "
|
|
f"tiling={enable_tiling}, tile={tile_size}"
|
|
)
|
|
else:
|
|
stagnation_count = 0
|
|
last_frames_processed = frames_processed
|
|
|
|
end_idx = min(start_idx + current_batch, total_frames)
|
|
ctx_start = max(0, start_idx - overlap_frames)
|
|
ctx_end = min(total_frames, end_idx + overlap_frames)
|
|
|
|
if ctx_end <= ctx_start:
|
|
raise RuntimeError(f"Empty chunk (ctx {ctx_start}→{ctx_end})")
|
|
|
|
latent_chunk = latents[:, :, ctx_start:ctx_end, :, :]
|
|
|
|
# Logging
|
|
should_log = verbose and (chunk_count % log_every_n == 0 or end_idx == total_frames)
|
|
if should_log:
|
|
mem_gb = torch.cuda.memory_allocated() / (1024**3) if torch.cuda.is_available() else 0
|
|
logger.info(
|
|
f"Chunk {chunk_count+1}/{est_total_chunks}: "
|
|
f"latent {start_idx}→{end_idx} | ctx {ctx_start}→{ctx_end} | "
|
|
f"VRAM: {mem_gb:.2f}GB"
|
|
)
|
|
|
|
# Decode with recovery
|
|
try:
|
|
with torch.no_grad():
|
|
if enable_tiling and hasattr(vae, 'decode_tiled'):
|
|
decoded_chunk = vae.decode_tiled(latent_chunk, tile_x=tile_size, tile_y=tile_size)
|
|
else:
|
|
decoded_chunk = vae.decode(latent_chunk)
|
|
decoded_chunk = decoded_chunk.cpu()
|
|
|
|
# Success - reset OOM counter
|
|
oom_retry_count = 0
|
|
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e).lower():
|
|
oom_retry_count += 1
|
|
|
|
if verbose:
|
|
logger.warning(f"OOM at frame {start_idx} (retry {oom_retry_count}/{self.MAX_OOM_RETRIES})")
|
|
|
|
# Cleanup
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize()
|
|
gc.collect()
|
|
|
|
# Stage 1: Enable tiling
|
|
if not enable_tiling:
|
|
logger.info("→ Stage 1: Enabling tiling")
|
|
enable_tiling = True
|
|
output_h, output_w = self.detect_output_size(
|
|
vae, latents, h_latent, w_latent, tile_size, verbose=False
|
|
)
|
|
continue
|
|
|
|
# Stage 2: Reduce batch
|
|
if current_batch > 1:
|
|
old_batch = current_batch
|
|
current_batch = max(1, current_batch // 2)
|
|
overlap_frames = min(initial_overlap, current_batch - 1)
|
|
logger.info(f"→ Stage 2: Batch {old_batch} → {current_batch}")
|
|
continue
|
|
|
|
# Stage 3: Reduce tile
|
|
if tile_size > 256:
|
|
old_tile = tile_size
|
|
tile_size = max(256, tile_size - 128)
|
|
logger.info(f"→ Stage 3: Tile {old_tile} → {tile_size}px")
|
|
continue
|
|
|
|
# All strategies exhausted, increment counter and continue
|
|
# Will hit MAX_OOM_RETRIES at top of loop
|
|
continue
|
|
else:
|
|
raise
|
|
|
|
# Normalize
|
|
decoded_chunk = self._normalize_output(decoded_chunk, aspect_ratio)
|
|
|
|
# Temporal trimming
|
|
front_trim = (start_idx - ctx_start) * time_scale
|
|
|
|
if end_idx == total_frames:
|
|
valid_frames = decoded_chunk[front_trim:]
|
|
else:
|
|
core_length = (end_idx - start_idx) * time_scale
|
|
valid_frames = decoded_chunk[front_trim:front_trim + core_length]
|
|
|
|
# Validation
|
|
actual = valid_frames.shape[0]
|
|
if end_idx < total_frames:
|
|
expected = (end_idx - start_idx) * time_scale
|
|
else:
|
|
expected = expected_frames - cumulative_output_frames
|
|
|
|
if should_log and abs(actual - expected) > 0:
|
|
logger.info(f" Chunk frames: expected {expected}, got {actual}")
|
|
|
|
# Spatial alignment
|
|
valid_frames = self._center_crop_to_reference(valid_frames, output_h, output_w)
|
|
|
|
# First chunk orientation check
|
|
if verbose and chunk_count == 0:
|
|
logger.info(f"✓ First chunk: shape {valid_frames.shape}")
|
|
if (aspect_ratio > 1.0 and valid_frames.shape[2] > valid_frames.shape[1]) or \
|
|
(aspect_ratio < 1.0 and valid_frames.shape[1] > valid_frames.shape[2]):
|
|
warnings.warn("Orientation mismatch! Output may be rotated.")
|
|
|
|
# Store
|
|
if use_disk_offload:
|
|
temp_path = os.path.join(self.temp_dir, f"chunk_{chunk_count:06d}.pt")
|
|
torch.save(valid_frames, temp_path)
|
|
temp_files.append(temp_path)
|
|
else:
|
|
final_output[current_write_idx:current_write_idx + actual] = valid_frames
|
|
current_write_idx += actual
|
|
|
|
# Progress update
|
|
processed_this_chunk = end_idx - start_idx
|
|
frames_processed += processed_this_chunk
|
|
cumulative_output_frames += actual
|
|
pbar.update(processed_this_chunk)
|
|
start_idx = end_idx
|
|
chunk_count += 1
|
|
|
|
# Cleanup
|
|
del latent_chunk, decoded_chunk, valid_frames
|
|
|
|
if chunk_count % 3 == 0 or current_batch <= 4:
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
torch.cuda.empty_cache()
|
|
|
|
# ======== FINAL ASSEMBLY ========
|
|
if use_disk_offload:
|
|
if verbose:
|
|
logger.info(f"🛡️ Assembling {len(temp_files)} chunks from disk...")
|
|
|
|
chunks = []
|
|
for i, temp_path in enumerate(temp_files):
|
|
chunk = torch.load(temp_path, map_location='cpu')
|
|
chunks.append(chunk)
|
|
|
|
# Cleanup
|
|
try:
|
|
os.remove(temp_path)
|
|
except Exception as e:
|
|
logger.debug(f"Failed to remove temp file: {e}")
|
|
|
|
# Merge every N chunks
|
|
if len(chunks) >= self.DISK_MERGE_CHUNK_SIZE or i == len(temp_files) - 1:
|
|
merged = torch.cat(chunks, dim=0)
|
|
chunks = [merged]
|
|
gc.collect()
|
|
|
|
final_output = chunks[0] if chunks else torch.empty((0, output_h, output_w, 3), dtype=torch.float32)
|
|
|
|
# Cleanup temp dir
|
|
self._cleanup_temp_dir()
|
|
self.temp_dir = None
|
|
|
|
# Trim if over-allocated
|
|
if final_output.shape[0] > expected_frames:
|
|
final_output = final_output[:expected_frames]
|
|
|
|
# Final cleanup
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
torch.cuda.empty_cache()
|
|
|
|
# Validation
|
|
actual_frames = final_output.shape[0]
|
|
|
|
if verbose:
|
|
logger.info(f"✅ Decode complete!")
|
|
logger.info(f" Output frames: {actual_frames}")
|
|
logger.info(f" Final shape: {final_output.shape}")
|
|
|
|
# Audio sync check
|
|
frame_diff = abs(actual_frames - expected_frames)
|
|
if frame_diff == 0:
|
|
if verbose:
|
|
logger.info(f" 🎵 Audio sync: PERFECT ✓")
|
|
elif frame_diff <= time_scale:
|
|
if verbose:
|
|
logger.info(f" ⚠️ Minor diff: {frame_diff} (rounding, OK)")
|
|
else:
|
|
if verbose:
|
|
logger.warning(f"Audio sync warning:")
|
|
logger.warning(f" Expected: {expected_frames}, Got: {actual_frames}")
|
|
logger.warning(f" Diff: {frame_diff} frames")
|
|
logger.warning(f" Try: force_time_scale={time_scale-1} or {time_scale+1}")
|
|
|
|
return (final_output,)
|
|
|
|
|
|
# ======== NODE REGISTRATION ========
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UniversalSmartVAEDecode": UniversalSmartVAEDecode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"UniversalSmartVAEDecode": "🎬 Universal VAE Decode (v11.2 Hardened)",
|
|
}
|