1275 lines
50 KiB
Python
1275 lines
50 KiB
Python
"""
|
||
Universal Smart VAE Video Decode - Streaming Node (FINAL with Ignore Warnings)
|
||
Production-ready standalone node for streaming video decoding.
|
||
|
||
Based on: UniversalSmartVAEDecode v11.2 + streaming extensions
|
||
Refinements: Claude (final polish + ignore warnings), Grok (implementation), Gemini (recovery), AI ensemble
|
||
|
||
Features:
|
||
- Frame-by-frame decoding (80% less RAM than standard pipeline)
|
||
- Crash recovery with resume capability
|
||
- Thumbnail previews in ComfyUI UI
|
||
- Multiple codecs: h264, h265, prores, FFV1 lossless
|
||
- Progress tracking with file write percentage
|
||
- Audio muxing support (now with AUDIO input support)
|
||
- Auto-detection of RAM/VRAM availability
|
||
- NEW: Ignore warnings mode (risk of black frames)
|
||
|
||
Version: 1.0.9 (Ignore Warnings + NaN handling from Claude)
|
||
License: MIT
|
||
"""
|
||
|
||
import torch
|
||
import os
|
||
import time
|
||
import json
|
||
import logging
|
||
from typing import Optional, Tuple
|
||
import gc
|
||
import comfy.utils
|
||
from comfy.model_management import throw_exception_if_processing_interrupted
|
||
import warnings
|
||
import folder_paths
|
||
import subprocess
|
||
import numpy as np
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Dependencies check
|
||
try:
|
||
import imageio
|
||
IMAGEIO_AVAILABLE = True
|
||
except ImportError:
|
||
IMAGEIO_AVAILABLE = False
|
||
warnings.warn(
|
||
"imageio not available. Install with: pip install imageio imageio-ffmpeg\n"
|
||
"This node requires imageio for video encoding."
|
||
)
|
||
|
||
try:
|
||
import numpy as np
|
||
NUMPY_AVAILABLE = True
|
||
except ImportError:
|
||
NUMPY_AVAILABLE = False
|
||
warnings.warn("numpy not available - required for this node")
|
||
|
||
try:
|
||
import cv2
|
||
CV2_AVAILABLE = True
|
||
except ImportError:
|
||
CV2_AVAILABLE = False
|
||
logger.info("cv2 not available - preview thumbnails will be skipped")
|
||
|
||
try:
|
||
import psutil
|
||
PSUTIL_AVAILABLE = True
|
||
except ImportError:
|
||
PSUTIL_AVAILABLE = False
|
||
logger.info("psutil not available - RAM estimation unavailable")
|
||
|
||
try:
|
||
import torchaudio
|
||
TORCHAUDIO_AVAILABLE = True
|
||
except ImportError:
|
||
TORCHAUDIO_AVAILABLE = False
|
||
logger.info("torchaudio not available - fallback to FFmpeg for AUDIO input.")
|
||
|
||
|
||
class StreamingVideoConfig:
|
||
"""Configuration for streaming video decode."""
|
||
|
||
CODECS = {
|
||
"h264": {
|
||
"codec": "libx264",
|
||
"ext": "mp4",
|
||
"output_params": ['-crf', '23', '-preset', 'slow'],
|
||
"pixel_format": "yuv420p",
|
||
"description": "H.264 (best compatibility, good quality)"
|
||
},
|
||
"h265": {
|
||
"codec": "libx265",
|
||
"ext": "mp4",
|
||
"output_params": ['-crf', '28', '-preset', 'slow'],
|
||
"pixel_format": "yuv420p",
|
||
"description": "H.265/HEVC (50% smaller files, slower encode)"
|
||
},
|
||
"prores": {
|
||
"codec": "prores_ks",
|
||
"ext": "mov",
|
||
"output_params": ['-profile:v', '3'],
|
||
"pixel_format": "yuv422p10le",
|
||
"description": "ProRes 422 (professional, large files)"
|
||
},
|
||
"ffv1": {
|
||
"codec": "ffv1",
|
||
"ext": "mkv",
|
||
"output_params": ['-level', '3', '-g', '1', '-slices', '4', '-slicecrc', '1'],
|
||
"pixel_format": "yuv444p",
|
||
"description": "FFV1 (lossless archival, very large files)"
|
||
}
|
||
}
|
||
|
||
PREVIEW_INTERVAL = 50
|
||
METADATA_SAVE_INTERVAL = 100
|
||
|
||
|
||
class StreamingVideoWriter:
|
||
"""Handles streaming video encoding with crash recovery."""
|
||
|
||
def __init__(self, output_path: str, fps: int, codec: str,
|
||
width: int, height: int, resume: bool = False):
|
||
self.output_path = output_path
|
||
self.fps = fps
|
||
self.codec = codec
|
||
self.width = width
|
||
self.height = height
|
||
self.resume = resume
|
||
|
||
self.metadata_path = output_path + ".metadata.json"
|
||
base, ext = os.path.splitext(output_path)
|
||
self.temp_path = f"{base}.tmp{ext}" if ext else output_path + ".tmp"
|
||
|
||
self.writer = None
|
||
self.frames_written = 0
|
||
self.session_frames_written = 0
|
||
self.last_preview = None
|
||
self.resume_prefix_path = None
|
||
|
||
if resume and os.path.exists(self.metadata_path):
|
||
self._load_metadata()
|
||
|
||
self._init_writer()
|
||
|
||
def _init_writer(self):
|
||
"""Initialize video writer."""
|
||
if not IMAGEIO_AVAILABLE:
|
||
raise RuntimeError(
|
||
"imageio required for video encoding.\n"
|
||
"Install: pip install imageio imageio-ffmpeg"
|
||
)
|
||
|
||
codec_config = StreamingVideoConfig.CODECS.get(
|
||
self.codec,
|
||
StreamingVideoConfig.CODECS["h264"]
|
||
)
|
||
|
||
write_path = self.temp_path
|
||
|
||
# Resume metadata tracks total decoded output frames, but encoded containers
|
||
# cannot be safely appended with imageio. Keep previous partial video as a
|
||
# prefix and encode only the continuation segment into `temp_path`.
|
||
if self.resume and self.frames_written > 0:
|
||
if os.path.exists(self.output_path):
|
||
self.resume_prefix_path = self.output_path
|
||
elif os.path.exists(self.temp_path):
|
||
temp_base, temp_ext = os.path.splitext(self.temp_path)
|
||
prefix_backup = f"{temp_base}.resume_prefix{temp_ext}" if temp_ext else self.temp_path + ".resume_prefix"
|
||
try:
|
||
os.replace(self.temp_path, prefix_backup)
|
||
self.resume_prefix_path = prefix_backup
|
||
except Exception as e:
|
||
logger.warning(f"Failed to preserve temp resume prefix: {e}")
|
||
self.frames_written = 0
|
||
else:
|
||
logger.warning(
|
||
"Resume metadata found but no partial video file exists. "
|
||
"Restarting decode from frame 0."
|
||
)
|
||
self.frames_written = 0
|
||
|
||
try:
|
||
self.writer = imageio.get_writer(
|
||
write_path,
|
||
fps=self.fps,
|
||
codec=codec_config["codec"],
|
||
quality=None,
|
||
pixelformat=codec_config["pixel_format"],
|
||
output_params=codec_config["output_params"],
|
||
macro_block_size=1
|
||
)
|
||
|
||
logger.info(f"Video writer initialized: {codec_config['description']}")
|
||
logger.info(f"Output: {write_path} ({self.width}x{self.height} @ {self.fps}fps)")
|
||
if self.resume_prefix_path:
|
||
logger.info(f"Resume prefix: {self.resume_prefix_path}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to initialize video writer: {e}")
|
||
raise
|
||
|
||
def write_frame(self, frame: torch.Tensor) -> Optional[np.ndarray]:
|
||
if not NUMPY_AVAILABLE:
|
||
raise RuntimeError("numpy required")
|
||
|
||
frame_np = (frame * 255.0).clamp(0, 255).byte().cpu().numpy()
|
||
|
||
self.writer.append_data(frame_np)
|
||
self.frames_written += 1
|
||
self.session_frames_written += 1
|
||
|
||
preview = None
|
||
if self.frames_written % StreamingVideoConfig.PREVIEW_INTERVAL == 0:
|
||
scale = min(1.0, 512.0 / max(self.width, self.height))
|
||
|
||
if scale < 1.0:
|
||
if CV2_AVAILABLE:
|
||
new_w = int(self.width * scale)
|
||
new_h = int(self.height * scale)
|
||
preview = cv2.resize(frame_np, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||
else:
|
||
preview = frame_np.copy()
|
||
|
||
self.last_preview = preview
|
||
|
||
if self.frames_written % StreamingVideoConfig.METADATA_SAVE_INTERVAL == 0:
|
||
self._save_metadata()
|
||
|
||
return preview
|
||
|
||
def _save_metadata(self):
|
||
metadata = {
|
||
"frames_written": self.frames_written,
|
||
"output_path": self.output_path,
|
||
"temp_path": self.temp_path,
|
||
"fps": self.fps,
|
||
"codec": self.codec,
|
||
"resolution": [self.width, self.height],
|
||
"timestamp": time.time()
|
||
}
|
||
|
||
try:
|
||
with open(self.metadata_path, 'w') as f:
|
||
json.dump(metadata, f, indent=2)
|
||
except Exception as e:
|
||
logger.debug(f"Metadata save failed: {e}")
|
||
|
||
def _load_metadata(self):
|
||
try:
|
||
with open(self.metadata_path, 'r') as f:
|
||
metadata = json.load(f)
|
||
|
||
self.frames_written = metadata.get("frames_written", 0)
|
||
logger.info(f"Resuming from frame {self.frames_written}")
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Metadata load failed: {e}")
|
||
self.frames_written = 0
|
||
|
||
@staticmethod
|
||
def _escape_concat_path(path: str) -> str:
|
||
return path.replace("'", "'\\''")
|
||
|
||
def _concat_videos(self, first_video: str, second_video: str) -> Optional[str]:
|
||
if not (os.path.exists(first_video) and os.path.exists(second_video)):
|
||
return None
|
||
|
||
try:
|
||
import imageio_ffmpeg
|
||
ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
|
||
except ImportError:
|
||
ffmpeg_exe = "ffmpeg"
|
||
|
||
base, ext = os.path.splitext(self.output_path)
|
||
concat_output = f"{base}.concat{ext}"
|
||
concat_list = f"{base}.concat.txt"
|
||
|
||
try:
|
||
with open(concat_list, "w", encoding="utf-8") as f:
|
||
f.write(f"file '{self._escape_concat_path(first_video)}'\n")
|
||
f.write(f"file '{self._escape_concat_path(second_video)}'\n")
|
||
|
||
cmd = [
|
||
ffmpeg_exe, "-y",
|
||
"-f", "concat",
|
||
"-safe", "0",
|
||
"-i", concat_list,
|
||
"-c", "copy",
|
||
concat_output,
|
||
]
|
||
subprocess.run(cmd, check=True, capture_output=True, timeout=600, text=True)
|
||
|
||
if not os.path.exists(concat_output):
|
||
logger.error("FFmpeg concat succeeded but output not created")
|
||
return None
|
||
|
||
os.replace(concat_output, self.output_path)
|
||
return self.output_path
|
||
|
||
except Exception as e:
|
||
logger.error(f"FFmpeg concat error: {e}")
|
||
return None
|
||
|
||
finally:
|
||
for cleanup_path in (concat_output, concat_list):
|
||
if os.path.exists(cleanup_path):
|
||
try:
|
||
os.remove(cleanup_path)
|
||
except Exception:
|
||
pass
|
||
|
||
def finalize(self, audio_path: Optional[str] = None) -> str:
|
||
if self.writer:
|
||
self.writer.close()
|
||
self.writer = None
|
||
|
||
if os.path.exists(self.metadata_path):
|
||
try:
|
||
os.remove(self.metadata_path)
|
||
except:
|
||
pass
|
||
|
||
temp_exists = os.path.exists(self.temp_path)
|
||
base_video_path = None
|
||
|
||
if self.resume_prefix_path and os.path.exists(self.resume_prefix_path):
|
||
if self.session_frames_written > 0 and temp_exists:
|
||
merged = self._concat_videos(self.resume_prefix_path, self.temp_path)
|
||
if merged:
|
||
base_video_path = merged
|
||
else:
|
||
logger.warning("Resume merge failed. Keeping previous partial output.")
|
||
base_video_path = self.resume_prefix_path
|
||
try:
|
||
os.remove(self.temp_path)
|
||
except Exception:
|
||
pass
|
||
else:
|
||
base_video_path = self.resume_prefix_path
|
||
if temp_exists:
|
||
try:
|
||
os.remove(self.temp_path)
|
||
except Exception:
|
||
pass
|
||
elif temp_exists:
|
||
try:
|
||
os.replace(self.temp_path, self.output_path)
|
||
base_video_path = self.output_path
|
||
except Exception as e:
|
||
logger.error(f"Failed to move temp file: {e}")
|
||
base_video_path = self.temp_path
|
||
else:
|
||
base_video_path = self.output_path
|
||
|
||
if base_video_path != self.output_path and os.path.exists(base_video_path):
|
||
try:
|
||
os.replace(base_video_path, self.output_path)
|
||
base_video_path = self.output_path
|
||
except Exception as e:
|
||
logger.warning(f"Could not normalize output path: {e}")
|
||
|
||
if audio_path and os.path.exists(audio_path):
|
||
final_path = self._mux_audio(base_video_path, audio_path)
|
||
if final_path == base_video_path and base_video_path == self.temp_path:
|
||
try:
|
||
os.replace(self.temp_path, self.output_path)
|
||
except:
|
||
pass
|
||
return self.output_path
|
||
return final_path
|
||
|
||
return base_video_path
|
||
|
||
def _mux_audio(self, video_path: str, audio_path: str) -> str:
|
||
if not os.path.exists(video_path):
|
||
logger.error(f"Video file not found: {video_path}")
|
||
return video_path
|
||
|
||
if not os.path.exists(audio_path):
|
||
logger.error(f"Audio file not found: {audio_path}")
|
||
return video_path
|
||
|
||
try:
|
||
import imageio_ffmpeg
|
||
ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
|
||
except ImportError:
|
||
ffmpeg_exe = "ffmpeg"
|
||
|
||
base, ext = os.path.splitext(self.output_path)
|
||
output_with_audio = f"{base}_audio{ext}"
|
||
|
||
codec_config = StreamingVideoConfig.CODECS.get(self.codec, StreamingVideoConfig.CODECS["h264"])
|
||
|
||
try:
|
||
cmd = [
|
||
ffmpeg_exe, "-y",
|
||
"-i", video_path,
|
||
"-i", audio_path,
|
||
"-map", "0:v:0",
|
||
"-map", "1:a:0",
|
||
"-c:v", codec_config["codec"],
|
||
]
|
||
|
||
cmd.extend(codec_config["output_params"])
|
||
cmd.extend([
|
||
"-pix_fmt", codec_config["pixel_format"],
|
||
"-c:a", "aac",
|
||
"-b:a", "192k",
|
||
"-shortest",
|
||
output_with_audio
|
||
])
|
||
|
||
logger.info(f"Running ffmpeg mux: {' '.join(cmd)}")
|
||
|
||
result = subprocess.run(
|
||
cmd,
|
||
check=True,
|
||
capture_output=True,
|
||
timeout=600,
|
||
text=True
|
||
)
|
||
|
||
if os.path.exists(output_with_audio):
|
||
try:
|
||
if video_path != self.output_path:
|
||
os.remove(video_path)
|
||
logger.info(f"Removed temp video: {video_path}")
|
||
except Exception as e:
|
||
logger.debug(f"Could not remove temp video: {e}")
|
||
|
||
logger.info(f"✓ Audio muxed: {output_with_audio}")
|
||
return output_with_audio
|
||
else:
|
||
logger.error("FFmpeg succeeded but output not created")
|
||
return video_path
|
||
|
||
except Exception as e:
|
||
logger.error(f"FFmpeg mux error: {e}")
|
||
return video_path
|
||
|
||
def __del__(self):
|
||
if self.writer:
|
||
try:
|
||
self.writer.close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
class SmartVAE_StreamingDecoder:
|
||
"""
|
||
Standalone streaming video decode node with warning ignore option.
|
||
"""
|
||
|
||
MAX_OOM_RETRIES = 5
|
||
|
||
@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 per decode 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 override (0=auto). E.g., 8 for LTX-Video."
|
||
}),
|
||
"enable_tiling": ("BOOLEAN", {
|
||
"default": False,
|
||
"tooltip": "Force tiling. Auto-enables on OOM."
|
||
}),
|
||
"tile_size": ("INT", {
|
||
"default": 512,
|
||
"min": 256,
|
||
"max": 2048,
|
||
"step": 64,
|
||
"tooltip": "Tile size in pixels."
|
||
}),
|
||
"verbose": ("BOOLEAN", {
|
||
"default": False,
|
||
"tooltip": "Show detailed logs."
|
||
}),
|
||
"video_output_path": ("STRING", {
|
||
"default": "",
|
||
"multiline": False,
|
||
"tooltip": "Output path. Empty = auto-generated."
|
||
}),
|
||
"fps": ("INT", {
|
||
"default": 24,
|
||
"min": 1,
|
||
"max": 120,
|
||
"step": 1,
|
||
"tooltip": "Frames per second"
|
||
}),
|
||
"codec": (list(StreamingVideoConfig.CODECS.keys()), {
|
||
"default": "h264",
|
||
"tooltip": "h264=compatible | h265=smaller | prores/ffv1=professional"
|
||
}),
|
||
"audio_path": ("STRING", {
|
||
"default": "",
|
||
"multiline": False,
|
||
"tooltip": "Optional audio file path to mux (alternative to AUDIO input)"
|
||
}),
|
||
"audio": ("AUDIO", {
|
||
"tooltip": "Optional AUDIO input (ComfyUI format). Priority over audio_path."
|
||
}),
|
||
"resume_on_crash": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Resume if crash occurred"
|
||
}),
|
||
"ignore_warnings": (["none", "minor", "all"], {
|
||
"default": "none",
|
||
"tooltip": (
|
||
"Warning handling mode:\n"
|
||
"• none = Stop on any issue (safest)\n"
|
||
"• minor = Continue if corruption <10% (partial black frames possible)\n"
|
||
"• all = Force decode anyway (high risk of black frames or crash)"
|
||
)
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "STRING")
|
||
RETURN_NAMES = ("preview_thumbs", "video_path")
|
||
FUNCTION = "decode"
|
||
OUTPUT_NODE = True
|
||
CATEGORY = "latent/video"
|
||
|
||
def __init__(self):
|
||
self._time_scale_cache = {}
|
||
self._force_scale_cache = {}
|
||
self._verbose = False
|
||
|
||
def _get_available_vram(self) -> Optional[float]:
|
||
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"VRAM detection failed: {e}")
|
||
return None
|
||
|
||
def _get_available_ram(self) -> Optional[float]:
|
||
if PSUTIL_AVAILABLE:
|
||
try:
|
||
return psutil.virtual_memory().available / (1024 ** 3)
|
||
except Exception as e:
|
||
logger.warning(f"RAM detection failed: {e}")
|
||
return None
|
||
|
||
def _estimate_chunk_vram(self, frames: int, channels: int, h: int, w: int,
|
||
time_scale: int = 1) -> float:
|
||
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:
|
||
bytes_per_frame = output_h * output_w * 3 * 4
|
||
return (expected_frames * bytes_per_frame) / (1024 ** 3)
|
||
|
||
def detect_time_scale(self, vae, latents: torch.Tensor, force_scale: int = 0,
|
||
verbose: bool = True) -> int:
|
||
vae_id = id(vae)
|
||
|
||
if force_scale > 0:
|
||
if self._force_scale_cache.get(vae_id) != force_scale:
|
||
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
|
||
|
||
self._force_scale_cache.pop(vae_id, None)
|
||
|
||
if vae_id in self._time_scale_cache:
|
||
return self._time_scale_cache[vae_id]
|
||
|
||
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 parse failed: {e}")
|
||
|
||
try:
|
||
total_latent_frames = latents.shape[2]
|
||
test_frames = min(5, total_latent_frames)
|
||
|
||
if test_frames <= 1:
|
||
if verbose:
|
||
logger.warning(f"Insufficient frames for detection. 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]
|
||
|
||
time_scale = max(1, (output_frames - 1) // (test_frames - 1))
|
||
|
||
if verbose:
|
||
logger.info(f"🔍 Auto-detected: {time_scale}x ({test_frames}→{output_frames})")
|
||
|
||
self._time_scale_cache[vae_id] = time_scale
|
||
|
||
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"Detection failed: {e}. Fallback: 1x")
|
||
|
||
self._time_scale_cache[vae_id] = 1
|
||
return 1
|
||
|
||
def detect_output_size(self, vae, latents, h_latent, w_latent, tile_size, verbose):
|
||
test_latent = latents[:, :, :1, :, :]
|
||
|
||
try:
|
||
with torch.no_grad():
|
||
if hasattr(vae, 'decode_tiled'):
|
||
test_out = vae.decode_tiled(test_latent, tile_x=tile_size, tile_y=tile_size)
|
||
else:
|
||
test_out = vae.decode(test_latent)
|
||
|
||
test_out = self._normalize_output(test_out, aspect_ratio=None)
|
||
|
||
output_h = test_out.shape[1]
|
||
output_w = test_out.shape[2]
|
||
|
||
if verbose:
|
||
logger.info(f" Detected output: {output_h}×{output_w}")
|
||
|
||
return output_h, output_w
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Could not decode test frame ({e}), using estimated size")
|
||
|
||
output_h = h_latent * 8
|
||
output_w = w_latent * 8
|
||
|
||
if verbose:
|
||
logger.info(f" Estimated output: {output_h}×{output_w}")
|
||
|
||
return output_h, output_w
|
||
|
||
def _extract_tensor_output(self, tensor):
|
||
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)}")
|
||
return tensor
|
||
|
||
def _output_frames_to_latent_index(self, output_frames: int, time_scale: int, total_frames: int) -> int:
|
||
"""Convert decoded output frame count to latent frame index for resume."""
|
||
if output_frames <= 0:
|
||
return 0
|
||
|
||
if time_scale <= 1:
|
||
return min(output_frames, total_frames)
|
||
|
||
latent_index = 1 + max(0, output_frames - 1) // time_scale
|
||
return min(latent_index, total_frames)
|
||
|
||
def _normalize_output(self, tensor, aspect_ratio: Optional[float] = None) -> torch.Tensor:
|
||
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)}")
|
||
|
||
if tensor.dtype != torch.float32:
|
||
tensor = tensor.float()
|
||
|
||
dim = tensor.dim()
|
||
|
||
if dim == 4:
|
||
if tensor.shape[1] in (3, 4):
|
||
tensor = tensor.permute(0, 2, 3, 1)
|
||
|
||
elif dim == 5:
|
||
shape = list(tensor.shape)
|
||
|
||
if shape[0] == 1:
|
||
tensor = tensor.squeeze(0)
|
||
shape = list(tensor.shape)
|
||
|
||
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 in {shape}")
|
||
|
||
remaining_idxs = [i for i in range(4) if i != c_idx]
|
||
remaining_sizes = [shape[i] for i in remaining_idxs]
|
||
|
||
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]
|
||
|
||
if aspect_ratio is not None:
|
||
if aspect_ratio > 1.0:
|
||
h_idx = spatial_large_idx
|
||
w_idx = spatial_small_idx
|
||
else:
|
||
w_idx = spatial_large_idx
|
||
h_idx = spatial_small_idx
|
||
else:
|
||
h_idx = spatial_large_idx
|
||
w_idx = spatial_small_idx
|
||
|
||
perm = [f_idx, h_idx, w_idx, c_idx]
|
||
tensor = tensor.permute(*perm)
|
||
|
||
else:
|
||
raise ValueError(f"Unsupported: {dim}D, shape {tensor.shape}")
|
||
|
||
min_val = tensor.min().item()
|
||
max_val = tensor.max().item()
|
||
|
||
if min_val < 0.0:
|
||
tensor = (tensor + 1.0) / 2.0
|
||
|
||
tensor = torch.nan_to_num(tensor, nan=0.0, posinf=1.0, neginf=0.0)
|
||
tensor = torch.clamp(tensor, 0.0, 1.0)
|
||
|
||
if self._verbose:
|
||
if min_val < 0.0 or max_val > 1.0:
|
||
logger.warning(f"Normalize: auto-scaled and clamped [{min_val:.4f}, {max_val:.4f}] → [0,1]")
|
||
if torch.isnan(tensor).any() or torch.isinf(tensor).any():
|
||
logger.warning("NaN/Inf detected and cleaned in output!")
|
||
|
||
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:
|
||
_, 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=False,
|
||
video_output_path="", fps=24, codec="h264", audio_path="", audio=None,
|
||
resume_on_crash=True, ignore_warnings="none"):
|
||
self._verbose = verbose
|
||
|
||
if not IMAGEIO_AVAILABLE or not NUMPY_AVAILABLE:
|
||
raise RuntimeError(
|
||
"Required dependencies missing.\n"
|
||
"Install: pip install imageio imageio-ffmpeg numpy"
|
||
)
|
||
|
||
latents = samples["samples"]
|
||
|
||
if latents.dim() == 4:
|
||
raise ValueError(
|
||
"This node is for video latents only (5D tensors).\n"
|
||
"For images (4D), use standard VAE Decode node."
|
||
)
|
||
|
||
batch, latent_channels, total_frames, h_latent, w_latent = latents.shape
|
||
|
||
if total_frames <= 0:
|
||
raise ValueError("Latent has no frames")
|
||
|
||
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)
|
||
|
||
output_h, output_w = self.detect_output_size(vae, latents, h_latent, w_latent, tile_size, verbose)
|
||
|
||
est_ram = self._estimate_output_ram(expected_frames, output_h, output_w)
|
||
if verbose:
|
||
logger.info(f" Estimated full-RAM usage (if not streaming): {est_ram:.2f}GB")
|
||
|
||
if not video_output_path:
|
||
output_dir = folder_paths.get_output_directory()
|
||
timestamp = int(time.time())
|
||
codec_ext = StreamingVideoConfig.CODECS[codec]["ext"]
|
||
video_output_path = os.path.join(output_dir, f"video_{timestamp}.{codec_ext}")
|
||
|
||
os.makedirs(os.path.dirname(video_output_path) or ".", exist_ok=True)
|
||
|
||
temp_audio_path = None
|
||
final_audio_path = audio_path
|
||
|
||
if audio is not None:
|
||
try:
|
||
if not isinstance(audio, dict) or 'waveform' not in audio or 'sample_rate' not in audio:
|
||
logger.warning(f"Invalid AUDIO format. Got: {type(audio)}")
|
||
else:
|
||
waveform = audio['waveform']
|
||
sample_rate = audio['sample_rate']
|
||
|
||
if waveform is None or sample_rate is None:
|
||
logger.error("AUDIO input missing 'waveform' or 'sample_rate'")
|
||
elif waveform.numel() == 0:
|
||
logger.error("AUDIO input waveform is empty")
|
||
else:
|
||
if isinstance(waveform, torch.Tensor):
|
||
waveform = waveform.cpu()
|
||
|
||
if waveform.dim() == 1:
|
||
waveform = waveform.unsqueeze(0)
|
||
elif waveform.dim() == 3:
|
||
waveform = waveform.squeeze(0)
|
||
|
||
timestamp = int(time.time() * 1000)
|
||
temp_audio_path = os.path.join(
|
||
folder_paths.get_temp_directory(),
|
||
f"temp_audio_{timestamp}_{os.getpid()}.wav"
|
||
)
|
||
|
||
try:
|
||
import imageio_ffmpeg
|
||
ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
|
||
except ImportError:
|
||
ffmpeg_exe = "ffmpeg"
|
||
|
||
waveform_np = (waveform.numpy() * 32767).astype('int16')
|
||
audio_channels = waveform_np.shape[0]
|
||
|
||
cmd = [
|
||
ffmpeg_exe, "-y",
|
||
"-f", "s16le",
|
||
"-ar", str(sample_rate),
|
||
"-ac", str(audio_channels),
|
||
"-i", "pipe:0",
|
||
"-c:a", "pcm_s16le",
|
||
temp_audio_path
|
||
]
|
||
|
||
process = subprocess.Popen(
|
||
cmd,
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE
|
||
)
|
||
|
||
stdout, stderr = process.communicate(input=waveform_np.T.tobytes())
|
||
|
||
if process.returncode != 0:
|
||
raise RuntimeError(f"FFmpeg failed: {stderr.decode()}")
|
||
|
||
final_audio_path = temp_audio_path
|
||
|
||
if verbose:
|
||
logger.info(f"✓ AUDIO saved (ffmpeg): {temp_audio_path}")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to process AUDIO input: {e}")
|
||
temp_audio_path = None
|
||
|
||
if final_audio_path and os.path.exists(final_audio_path):
|
||
logger.info(f"🔊 Audio source ready: {final_audio_path}")
|
||
elif final_audio_path:
|
||
logger.warning(f"Audio path specified but file not found: {final_audio_path}")
|
||
final_audio_path = ""
|
||
|
||
logger.info("🎬 STREAMING VIDEO DECODE")
|
||
logger.info(f" Output: {video_output_path}")
|
||
logger.info(f" Codec: {codec} @ {fps}fps")
|
||
logger.info(f" Resolution: {output_h}x{output_w}")
|
||
logger.info(f" Expected frames: ~{expected_frames}")
|
||
if final_audio_path:
|
||
logger.info(f" Audio source: {final_audio_path}")
|
||
if ignore_warnings != "none":
|
||
logger.warning(f"⚠️ Ignore warnings mode: {ignore_warnings} (risk of black frames)")
|
||
|
||
try:
|
||
return self._streaming_decode(
|
||
vae, latents, video_output_path, fps, codec, final_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
|
||
)
|
||
finally:
|
||
if temp_audio_path and os.path.exists(temp_audio_path):
|
||
for attempt in range(5):
|
||
try:
|
||
time.sleep(0.5)
|
||
os.remove(temp_audio_path)
|
||
if verbose:
|
||
logger.info(f"✓ Cleaned temp audio: {temp_audio_path}")
|
||
break
|
||
except Exception as e:
|
||
if attempt == 4:
|
||
logger.warning(f"Could not remove temp audio: {e}")
|
||
continue
|
||
|
||
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"):
|
||
"""Streaming decode implementation with warning ignore."""
|
||
|
||
batch, _, total_frames, _, _ = latents.shape
|
||
|
||
if verbose:
|
||
logger.info("🔍 Validating latent before decode...")
|
||
|
||
try:
|
||
_ = latents[0, 0, 0, 0, 0].item()
|
||
|
||
if torch.isnan(latents).any():
|
||
nan_count = torch.isnan(latents).sum().item()
|
||
total_elements = latents.numel()
|
||
nan_percent = (nan_count / total_elements) * 100
|
||
|
||
logger.error("=" * 70)
|
||
logger.error("🚨 CORRUPTED LATENT DETECTED!")
|
||
logger.error(f" NaN values: {nan_count:,} / {total_elements:,} ({nan_percent:.2f}%)")
|
||
logger.error(f" Shape: {latents.shape}")
|
||
|
||
nan_frames = torch.isnan(latents).any(dim=(0,1,3,4)).nonzero(as_tuple=True)[0]
|
||
if len(nan_frames) > 0:
|
||
first_bad = nan_frames[0].item()
|
||
last_bad = nan_frames[-1].item()
|
||
logger.error(f" Affected frames: {first_bad} to {last_bad}")
|
||
logger.error(f" → Corruption likely started around frame {first_bad}")
|
||
|
||
logger.error("")
|
||
logger.error("💊 RECOMMENDED FIXES:")
|
||
logger.error(" 1. REDUCE VIDEO LENGTH (most effective)")
|
||
logger.error(f" Current: {total_frames} frames → try {int(total_frames * 0.6)}")
|
||
logger.error(" 2. REDUCE RESOLUTION")
|
||
logger.error(" 3. Lower CFG scale (e.g. 3.5 instead of 7.0)")
|
||
logger.error(" 4. Enable CPU offload or tiled sampling")
|
||
logger.error(" 5. Clear VRAM before workflow")
|
||
logger.error("=" * 70)
|
||
|
||
if ignore_warnings == "none":
|
||
raise ValueError(
|
||
f"Latent contains {nan_percent:.1f}% NaN values - cannot decode safely.\n"
|
||
"To force decode anyway (may produce black frames):\n"
|
||
" Set 'ignore_warnings' to 'minor' (<10% corruption) or 'all' (high risk)"
|
||
)
|
||
|
||
elif ignore_warnings == "minor":
|
||
if nan_percent > 10.0:
|
||
raise ValueError(
|
||
f"Latent is {nan_percent:.1f}% corrupted (limit: 10% for 'minor' mode).\n"
|
||
"This is too severe - output would be mostly black.\n\n"
|
||
"Options:\n"
|
||
" 1. Reduce video length/resolution and re-sample\n"
|
||
" 2. Set 'ignore_warnings' to 'all' (not recommended)"
|
||
)
|
||
|
||
logger.warning("=" * 70)
|
||
logger.warning("⚠️ CONTINUING WITH PARTIALLY CORRUPTED LATENT (user override)")
|
||
logger.warning(f" Corruption level: {nan_percent:.2f}% (under 10% threshold)")
|
||
logger.warning(f" Affected frames {first_bad}-{last_bad} will likely be BLACK")
|
||
logger.warning(" The rest of the video may be OK")
|
||
logger.warning("=" * 70)
|
||
|
||
latents = torch.nan_to_num(latents, nan=0.0, posinf=1.0, neginf=0.0)
|
||
|
||
elif ignore_warnings == "all":
|
||
logger.error("=" * 70)
|
||
logger.error("🚨 FORCING DECODE WITH SEVERELY CORRUPTED LATENT!")
|
||
logger.error(f" Corruption: {nan_percent:.2f}%")
|
||
logger.error(" ⚠️ WARNING: Output will likely be:")
|
||
logger.error(" • Mostly/entirely BLACK")
|
||
logger.error(" • May CRASH during decode")
|
||
logger.error(" • Waste time and produce unusable file")
|
||
logger.error("")
|
||
logger.error(" You have been warned. Proceeding anyway...")
|
||
logger.error("=" * 70)
|
||
|
||
latents = torch.nan_to_num(latents, nan=0.0, posinf=1.0, neginf=0.0)
|
||
|
||
if torch.isinf(latents).any():
|
||
raise ValueError("❌ Latent contains Inf - upstream overflow!")
|
||
|
||
if latents.numel() == 0:
|
||
raise ValueError("❌ Latent is empty - upstream allocation failed!")
|
||
|
||
if verbose:
|
||
min_val = latents.min().item()
|
||
max_val = latents.max().item()
|
||
logger.info(f" ✓ Latent range: [{min_val:.4f}, {max_val:.4f}]")
|
||
logger.info(f" ✓ Shape: {latents.shape}, Device: {latents.device}")
|
||
|
||
except RuntimeError as e:
|
||
error_msg = str(e).lower()
|
||
if "cuda" in error_msg or "memory" in error_msg or "out of memory" in error_msg:
|
||
logger.error("=" * 70)
|
||
logger.error("🚨 SILENT OOM DETECTED IN UPSTREAM NODE (SAMPLER)!")
|
||
logger.error(f" Latent shape: {latents.shape}")
|
||
logger.error(f" Expected size: ~{(latents.numel() * 4 / 1024**3):.2f} GB")
|
||
logger.error(f" Device: {latents.device}")
|
||
logger.error(f" Error: {e}")
|
||
logger.error("=" * 70)
|
||
logger.error("💡 ROOT CAUSE: The sampler ran out of VRAM without throwing error.")
|
||
logger.error(" The latent is CORRUPTED and cannot be decoded properly.")
|
||
logger.error("")
|
||
logger.error("💊 SOLUTIONS:")
|
||
logger.error(" 1. REDUCE VIDEO LENGTH (most effective)")
|
||
logger.error(" 2. REDUCE RESOLUTION")
|
||
logger.error(" 3. Lower CFG scale")
|
||
logger.error(" 4. Enable CPU offload or tiled sampling")
|
||
logger.error(" 5. Clear VRAM before workflow")
|
||
logger.error("=" * 70)
|
||
|
||
logger.warning("⚠️ Attempting emergency CPU recovery...")
|
||
try:
|
||
latents = latents.cpu()
|
||
vae.first_stage_model.to('cpu')
|
||
logger.info(" ✓ Moved to CPU - decode will be SLOW but may work")
|
||
logger.info(" ⚠️ If you get black frames, the latent is too corrupted")
|
||
except Exception as recovery_error:
|
||
raise RuntimeError(
|
||
f"Latent is corrupted beyond recovery.\n"
|
||
f"Original error: {e}\n"
|
||
f"Recovery failed: {recovery_error}\n\n"
|
||
f"You MUST reduce video length or resolution."
|
||
) from e
|
||
else:
|
||
raise
|
||
|
||
if total_frames > 1000:
|
||
original_batch = frames_per_batch
|
||
|
||
if total_frames > 2000:
|
||
safe_batch = 4
|
||
elif total_frames > 1500:
|
||
safe_batch = 6
|
||
else:
|
||
safe_batch = 8
|
||
|
||
if frames_per_batch > safe_batch:
|
||
frames_per_batch = safe_batch
|
||
if verbose:
|
||
logger.info("=" * 70)
|
||
logger.info(f"🛡️ LONG SEQUENCE DETECTED: {total_frames} frames")
|
||
logger.info(f" Auto-reducing batch: {original_batch} → {frames_per_batch}")
|
||
logger.info(f" (Prevents NaN accumulation in VAE)")
|
||
logger.info("=" * 70)
|
||
|
||
if total_frames > 2500 and latents.device.type == 'cuda':
|
||
logger.warning("=" * 70)
|
||
logger.warning(f"🚨 MEGA-SEQUENCE: {total_frames} frames")
|
||
logger.warning(" Forcing FULL CPU DECODE to prevent crash")
|
||
logger.warning(" This will be VERY SLOW (~10x slower)")
|
||
logger.warning("=" * 70)
|
||
|
||
latents = latents.cpu()
|
||
vae.first_stage_model.to('cpu')
|
||
frames_per_batch = min(frames_per_batch, 4)
|
||
gc.collect()
|
||
torch.cuda.empty_cache()
|
||
|
||
if verbose:
|
||
logger.info("✓ Latent validation complete")
|
||
|
||
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))
|
||
initial_overlap = overlap_frames
|
||
|
||
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, latent_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"📉 VRAM reduction: {old_batch} → {frames_per_batch}")
|
||
|
||
if verbose:
|
||
logger.info(f" Batch: {frames_per_batch}, Overlap: {overlap_frames}")
|
||
|
||
current_batch = frames_per_batch
|
||
resume_output_frames = writer.frames_written if resume_on_crash else 0
|
||
start_idx = (
|
||
self._output_frames_to_latent_index(resume_output_frames, time_scale, total_frames)
|
||
if resume_on_crash else 0
|
||
)
|
||
if verbose and resume_output_frames > 0:
|
||
logger.info(
|
||
f"Resuming at latent frame {start_idx}/{total_frames} "
|
||
f"(decoded frames already written: {resume_output_frames})"
|
||
)
|
||
frames_processed = 0
|
||
last_frames_processed = -1
|
||
stagnation_count = 0
|
||
MAX_STAGNATION = 3
|
||
oom_retry_count = 0
|
||
|
||
pbar = comfy.utils.ProgressBar(total_frames)
|
||
preview_frames = []
|
||
|
||
while start_idx < total_frames:
|
||
throw_exception_if_processing_interrupted()
|
||
|
||
if oom_retry_count >= self.MAX_OOM_RETRIES:
|
||
raise RuntimeError(f"Exceeded {self.MAX_OOM_RETRIES} OOM retries")
|
||
|
||
if frames_processed == last_frames_processed:
|
||
stagnation_count += 1
|
||
if stagnation_count >= MAX_STAGNATION:
|
||
raise RuntimeError(f"Stalled at frame {start_idx}")
|
||
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")
|
||
|
||
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_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 {start_idx} (retry {oom_retry_count}/{self.MAX_OOM_RETRIES})")
|
||
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
torch.cuda.synchronize()
|
||
gc.collect()
|
||
|
||
if not enable_tiling:
|
||
enable_tiling = True
|
||
continue
|
||
if current_batch > 1:
|
||
current_batch = max(1, current_batch // 2)
|
||
overlap_frames = min(initial_overlap, current_batch - 1)
|
||
continue
|
||
if tile_size > 256:
|
||
tile_size = max(256, tile_size - 128)
|
||
continue
|
||
|
||
continue
|
||
else:
|
||
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_length = (end_idx - start_idx) * time_scale
|
||
valid_frames = decoded_chunk[front_trim:front_trim + core_length]
|
||
|
||
valid_frames = self._center_crop_to_reference(valid_frames, output_h, output_w)
|
||
|
||
for frame_idx in range(valid_frames.shape[0]):
|
||
frame = valid_frames[frame_idx]
|
||
|
||
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)
|
||
|
||
processed_this_chunk = end_idx - start_idx
|
||
frames_processed += processed_this_chunk
|
||
pbar.update(processed_this_chunk)
|
||
|
||
progress_pct = (writer.frames_written / expected_frames) * 100
|
||
if verbose:
|
||
logger.info(f"Progress: {writer.frames_written}/{expected_frames} ({progress_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 ignore_warnings != "none":
|
||
logger.warning("=" * 70)
|
||
logger.warning("⚠️ VIDEO CREATED WITH WARNINGS IGNORED")
|
||
logger.warning(f" Mode: '{ignore_warnings}'")
|
||
logger.warning(f" Some frames MAY BE BLACK or corrupted")
|
||
logger.warning(" Recommend: Check output before sharing")
|
||
logger.warning("=" * 70)
|
||
|
||
if verbose:
|
||
logger.info(f"✅ Streaming complete!")
|
||
logger.info(f" File: {final_path}")
|
||
logger.info(f" Frames: {writer.frames_written}")
|
||
file_size_mb = os.path.getsize(final_path) / (1024 * 1024)
|
||
logger.info(f" Size: {file_size_mb:.2f} MB")
|
||
|
||
if preview_frames:
|
||
preview_tensor = torch.stack(preview_frames, dim=0)
|
||
else:
|
||
preview_tensor = torch.empty(0, dtype=torch.float32)
|
||
|
||
logger.info("📸 NOTE: 'preview_thumbs' output is for live monitoring only (last 5 thumbnails). Full video is saved to disk.")
|
||
|
||
return (preview_tensor, final_path)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"SmartVAE_StreamingDecoder":SmartVAE_StreamingDecoder,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"SmartVAE_StreamingDecoder": "🎬 SmartVAE_StreamingDecoder (Streaming)",
|
||
}
|