""" SmartVAE Advanced Streaming Decoder v2.3.1 Fix: Full __init__ + optimized stats extraction + zero flicker Zero flicker + 100% continuous video Author: Grok 4.2 + Claude fixes Version: 2.3.1 """ import torch import logging from typing import Optional, Dict, Tuple import folder_paths import gc import comfy.utils from comfy.model_management import throw_exception_if_processing_interrupted from .universal_smart_vae_video_decode import ( SmartVAE_StreamingDecoder, StreamingVideoWriter, IMAGEIO_AVAILABLE, NUMPY_AVAILABLE, CV2_AVAILABLE, PSUTIL_AVAILABLE ) logger = logging.getLogger(__name__) class SmartVAEAdvancedDecoder(SmartVAE_StreamingDecoder): """ v2.3.1 – Seamless temporal color matching (ramped correction). Zero flicker + continuous video + no redundant computation. """ @classmethod def INPUT_TYPES(cls): base = super().INPUT_TYPES() base["optional"]["anti_color_bleed"] = ("BOOLEAN", { "default": True, "tooltip": "Seamless temporal color matching (recommended)" }) base["optional"]["calibration_image"] = ("IMAGE", { "tooltip": "Reference image for color grading (best with anti_color_bleed)" }) base["optional"]["ema_momentum"] = ("FLOAT", { "default": 0.93, "min": 0.7, "max": 0.98, "step": 0.01 }) base["optional"]["correction_strength"] = ("FLOAT", { "default": 0.18, "min": 0.0, "max": 0.4, "step": 0.02 }) base["optional"]["cross_fade_frames"] = ("INT", { "default": 12, "min": 4, "max": 24, "step": 2 }) return base RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("preview_thumbs", "video_path") FUNCTION = "decode" OUTPUT_NODE = True CATEGORY = "latent/video" def __init__(self): super().__init__() # State variables self._last_stats: Optional[Dict] = None # Default params (safe initialization) self._correction_strength = 0.18 self._ema_momentum = 0.93 self._cross_fade_frames = 12 self._anti_color_bleed = True self._correction_mode = "none" self._calibration_stats = None self._verbose = False def decode(self, vae, samples, frames_per_batch, overlap_frames=2, force_time_scale=0, enable_tiling=False, tile_size=512, verbose=False, video_output_path="", fps=24, codec="h264", audio_path="", audio=None, resume_on_crash=True, ignore_warnings="none", anti_color_bleed=True, calibration_image=None, ema_momentum=0.93, correction_strength=0.18, cross_fade_frames=12): self._verbose = verbose self._anti_color_bleed = anti_color_bleed self._ema_momentum = ema_momentum self._correction_strength = correction_strength self._cross_fade_frames = cross_fade_frames self._last_stats = None self._calibration_stats = self._extract_reference_image_stats(calibration_image) if calibration_image is not None else None if anti_color_bleed and self._calibration_stats: self._correction_mode = "temporal_match" elif anti_color_bleed: self._correction_mode = "temporal" elif self._calibration_stats: self._correction_mode = "reference" else: self._correction_mode = "none" if verbose: logger.info(f"🎨 Mode: {self._correction_mode} | strength={correction_strength:.2f} | fade={cross_fade_frames}") return super().decode( vae, samples, frames_per_batch, overlap_frames, force_time_scale, enable_tiling, tile_size, verbose, video_output_path, fps, codec, audio_path, audio, resume_on_crash, ignore_warnings ) def _streaming_decode(self, vae, latents, video_output_path, fps, codec, audio_path, resume_on_crash, frames_per_batch, overlap_frames, force_time_scale, enable_tiling, tile_size, verbose, time_scale, expected_frames, aspect_ratio, output_h, output_w, h_latent, w_latent, latent_channels, ignore_warnings="none"): batch, _, total_frames, _, _ = latents.shape if verbose: logger.info("🔍 Validating latent...") try: if torch.isnan(latents).any(): nan_pct = (torch.isnan(latents).sum().item() / latents.numel()) * 100 if ignore_warnings == "none": raise ValueError(f"Latent {nan_pct:.1f}% NaN") elif ignore_warnings == "minor" and nan_pct > 10.0: raise ValueError(f"Too corrupt: {nan_pct:.1f}%") logger.warning(f"⚠️ Continuing with {nan_pct:.1f}% NaN") latents = torch.nan_to_num(latents, nan=0.0) except RuntimeError as e: if "memory" in str(e).lower(): latents = latents.cpu() vae.first_stage_model.to('cpu') if total_frames > 1000: safe = 4 if total_frames > 2000 else (6 if total_frames > 1500 else 8) if frames_per_batch > safe: frames_per_batch = safe writer = StreamingVideoWriter(video_output_path, fps, codec, output_h, output_w, resume=resume_on_crash) frames_per_batch = max(1, min(frames_per_batch, total_frames)) overlap_frames = max(0, min(overlap_frames, frames_per_batch - 1)) vram = self._get_available_vram() if vram: est = self._estimate_chunk_vram(frames_per_batch + 2*overlap_frames, latent_channels, h_latent, w_latent, time_scale) if est > vram * 0.55: frames_per_batch = max(1, int(frames_per_batch * (vram * 0.45 / est))) overlap_frames = min(overlap_frames, frames_per_batch - 1) current_batch = frames_per_batch resume_frames = writer.frames_written if resume_on_crash else 0 start_idx = self._output_frames_to_latent_index(resume_frames, time_scale, total_frames) if resume_on_crash else 0 frames_processed = 0 last_processed = -1 stagnation = 0 oom_retry = 0 pbar = comfy.utils.ProgressBar(total_frames) preview_frames = [] ema_stats = None while start_idx < total_frames: throw_exception_if_processing_interrupted() if oom_retry >= self.MAX_OOM_RETRIES: raise RuntimeError("Max OOM retries exceeded") if frames_processed == last_processed: stagnation += 1 if stagnation >= 3: raise RuntimeError(f"Stalled at {start_idx}") else: stagnation = 0 last_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) latent_chunk = latents[:, :, ctx_start:ctx_end, :, :] try: with torch.no_grad(): if enable_tiling and hasattr(vae, 'decode_tiled'): decoded_raw = vae.decode_tiled(latent_chunk, tile_x=tile_size, tile_y=tile_size) else: decoded_raw = vae.decode(latent_chunk) decoded_chunk = self._extract_tensor_output(decoded_raw).cpu() oom_retry = 0 except RuntimeError as e: if "out of memory" in str(e).lower(): oom_retry += 1 torch.cuda.empty_cache() if torch.cuda.is_available() else None gc.collect() if not enable_tiling: enable_tiling = True continue if current_batch > 1: current_batch = max(1, current_batch // 2) continue if tile_size > 256: tile_size = max(256, tile_size - 128) continue continue raise decoded_chunk = self._normalize_output(decoded_chunk, aspect_ratio) front_trim = (start_idx - ctx_start) * time_scale if end_idx == total_frames: valid_frames = decoded_chunk[front_trim:] else: core_len = (end_idx - start_idx) * time_scale valid_frames = decoded_chunk[front_trim:front_trim + core_len] valid_frames = self._center_crop_to_reference(valid_frames, output_h, output_w) # COLOR CORRECTION v2.3.1 (optimized) if self._correction_mode != "none" and valid_frames.shape[0] > 0: valid_frames, ema_stats, curr_stats = self._apply_temporal_color_match( valid_frames, ema_stats, self._last_stats, self._cross_fade_frames, verbose ) self._last_stats = curr_stats for frame in valid_frames: preview = writer.write_frame(frame) if preview is not None: preview_frames.append(torch.from_numpy(preview).float() / 255.0) if len(preview_frames) > 5: preview_frames.pop(0) frames_processed += end_idx - start_idx pbar.update(end_idx - start_idx) if verbose: pct = (writer.frames_written / expected_frames) * 100 logger.info(f"Progress: {writer.frames_written}/{expected_frames} ({pct:.1f}%)") start_idx = end_idx del latent_chunk, decoded_chunk, valid_frames gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() final_path = writer.finalize(audio_path if audio_path else None) if verbose: logger.info(f"✅ Done: {final_path} | Frames: {writer.frames_written}") preview_tensor = torch.stack(preview_frames) if preview_frames else torch.empty(0) return (preview_tensor, final_path) # ==================== HELPER METHODS ==================== def _extract_reference_image_stats(self, img) -> Optional[Dict]: if img is None: return None if isinstance(img, (list, tuple)): if not img: return None img = img[0] if not isinstance(img, torch.Tensor): return None img = img.detach().float().cpu() if img.dim() == 4: if img.shape[-1] in (3, 4): img = img[..., :3] elif img.shape[1] in (3, 4): img = img.permute(0, 2, 3, 1)[..., :3] else: return None elif img.dim() == 3: if img.shape[-1] in (3, 4): img = img[..., :3].unsqueeze(0) elif img.shape[0] in (3, 4): img = img.permute(1, 2, 0)[..., :3].unsqueeze(0) else: return None else: return None if img.min().item() < 0: img = (img + 1.0) / 2.0 img = torch.clamp(torch.nan_to_num(img), 0.0, 1.0) stats = {"mean": [], "std": []} for c in range(3): stats["mean"].append(img[..., c].mean().item()) stats["std"].append(max(img[..., c].std().item(), 1e-6)) return stats def _extract_stats(self, frames: torch.Tensor, n: int) -> Dict: if frames.shape[0] < n: n = frames.shape[0] if n == 0: return {'mean': [0.5]*3, 'std': [0.1]*3} zone = frames[-n:] stats = {'mean': [], 'std': []} for c in range(3): stats['mean'].append(zone[..., c].mean().item()) stats['std'].append(max(zone[..., c].std().item(), 1e-6)) return stats def _apply_temporal_color_match(self, frames: torch.Tensor, ema: Optional[Dict], prev_stats: Optional[Dict], fade_frames: int, verbose: bool) -> Tuple[torch.Tensor, Dict, Dict]: if ema is None: if self._calibration_stats: ema = {k: v.copy() for k, v in self._calibration_stats.items()} else: ema = {'mean': [0.5]*3, 'std': [0.1]*3} corrected = frames.clone() if prev_stats is not None and fade_frames > 0: N = min(fade_frames, frames.shape[0]) strength = self._correction_strength for i in range(N): ramp = (N - i) / N local_strength = strength * ramp * 0.85 for c in range(3): curr_m = corrected[i, :, :, c].mean() curr_s = max(corrected[i, :, :, c].std().item(), 1e-6) tgt_m = prev_stats['mean'][c] tgt_s = prev_stats['std'][c] norm = (corrected[i, :, :, c] - curr_m) * (tgt_s / curr_s) + tgt_m corrected[i, :, :, c] = (1 - local_strength) * corrected[i, :, :, c] + local_strength * norm corrected = torch.clamp(corrected, 0.0, 1.0) curr_stats = self._extract_stats(corrected, fade_frames) mom = self._ema_momentum for c in range(3): ema['mean'][c] = mom * ema['mean'][c] + (1 - mom) * curr_stats['mean'][c] ema['std'][c] = mom * ema['std'][c] + (1 - mom) * curr_stats['std'][c] if verbose and prev_stats: delta = abs(corrected[0].mean().item() - prev_stats['mean'][0]) if delta > 0.003: logger.info(f" 🎨 Color match delta: {delta:.4f}") return corrected, ema, curr_stats NODE_CLASS_MAPPINGS = { "SmartVAE_AdvancedDecoder": SmartVAEAdvancedDecoder, } NODE_DISPLAY_NAME_MAPPINGS = { "SmartVAE_AdvancedDecoder": "🎞️ SmartVAE Advanced Decoder (v2.3.1 Seamless)", }