Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4490bd1f48 | ||
|
|
baec4b634f | ||
|
|
5a4bf428f3 | ||
|
|
43e70bf637 | ||
|
|
855f8b91b3 | ||
|
|
f561743054 | ||
|
|
396f323eae | ||
|
|
6226878411 | ||
|
|
2f8d2ccf9a | ||
|
|
b0f01f2d99 | ||
|
|
2214f3afde | ||
|
|
aeebd49f7f | ||
|
|
e178b72d89 | ||
|
|
8ad4c8fa4e | ||
|
|
241b632cfc | ||
|
|
27ed3333fd | ||
|
|
c6997fd9c2 | ||
|
|
0e849d20cd | ||
|
|
a65ddadc00 | ||
|
|
48bfbae05d | ||
|
|
04475bdc24 | ||
|
|
7bb936749f | ||
|
|
4fc3296c81 | ||
|
|
d69b65f7e4 | ||
|
|
15cb24089a | ||
|
|
c52280881a | ||
|
|
4b0b7d58b6 | ||
|
|
39d8d4bf19 | ||
|
|
f2f4916c05 | ||
|
|
f75bcc7f37 | ||
|
|
0c2a546c12 | ||
|
|
32f9900ecd | ||
|
|
84abef8de0 | ||
|
|
f3136dd20c | ||
|
|
93a6355517 | ||
|
|
5c07a92b33 | ||
|
|
d114e4958a |
@@ -36,6 +36,36 @@ We're actively working on improvements and new features. To stay informed:
|
||||
|
||||
## 🚀 Release Notes
|
||||
|
||||
**2025.12.24 - Version 2.5.24**
|
||||
|
||||
- **🍎 Fix: MPS memory leak regression** - Restored MPS cache clearing after VAE encode/decode operations that was accidentally removed during code cleanup in v2.5.23
|
||||
|
||||
**2025.12.24 - Version 2.5.23**
|
||||
|
||||
- **🔒 Security: Prevent code execution in model loading** - Added protection against malicious .pth files by restricting deserialization to tensors only
|
||||
- **🎥 Fix: FFmpeg video writer reliability** - Resolved ffmpeg process hanging issues by redirecting stderr and adding buffer flush, with improved error messages for debugging *(thanks [@thehhmdb](https://github.com/thehhmdb))*
|
||||
- **⚡ Fix: GGUF VAE model support** - Enabled automatic weight dequantization for convolution operations, making GGUF-quantized VAE models fully functional *(thanks [@naxci1](https://github.com/naxci1))*
|
||||
- **🛡️ Fix: VAE slicing edge cases** - Protected against division by zero crashes when using small split sizes with high temporal downsampling *(thanks [@naxci1](https://github.com/naxci1))*
|
||||
- **🎨 Fix: LAB color transfer precision** - Resolved dtype mismatch errors during video upscaling by ensuring consistent float types before matrix operations
|
||||
- **🔧 Fix: PyTorch 2.9+ compatibility** - Extended Conv3d memory workaround to all PyTorch 2.9+ versions, fixing 3x VRAM usage on newer PyTorch releases
|
||||
- **📦 Fix: Bitsandbytes compatibility** - Added ValueError exception handling for Intel Gaudi version detection failures on non-Gaudi systems
|
||||
- **🍎 MPS: Memory optimization** - Reduced memory usage during encode/decode operations on Apple Silicon *(thanks [@s-cerevisiae](https://github.com/s-cerevisiae))*
|
||||
|
||||
|
||||
**2025.12.13 - Version 2.5.22**
|
||||
|
||||
- **🎬 CLI: FFmpeg video backend with 10-bit support** - New `--video_backend ffmpeg` and `--10bit` flags enable x265 encoding with 10-bit color depth, reducing banding artifacts in gradients compared to 8-bit OpenCV output *(based on PR by [@thehhmdb](https://github.com/thehhmdb) - thank you!)*
|
||||
- **🍎 Fix: MPS bicubic upscaling compatibility** - Added CPU fallback for bicubic+antialias interpolation on PyTorch versions before 2.8.0, resolving RGBA alpha upscaling errors on Apple Silicon
|
||||
- **⚡ Fix: Cross-platform histogram matching** - Replaced scatter_ operation with argsort+index_select for improved reliability across CUDA, ROCm, and MPS backends
|
||||
- **🧹 MPS: Remove sync overhead** - Reverted unnecessary `torch.mps.synchronize()` calls introduced in v2.5.21 for consistent behavior with CUDA pipeline
|
||||
|
||||
**2025.12.12 - Version 2.5.21**
|
||||
|
||||
- **🛠️ Fix: GGUF dequantization error on MPS** - Resolved shape mismatch error introduced in 2.5.20 by skipping GGUF quantized buffers in precision conversion - these must remain in packed format for on-the-fly dequantization during inference
|
||||
- **🍎 MPS: Eliminate CPU sync overhead** - Skip unnecessary CPU tensor offload on Apple Silicon unified memory architecture, preventing sync stalls that caused slowdowns. Input images and output video now stay on MPS device throughout the pipeline
|
||||
- **⚡ MPS: Preload text embeddings** - Load text embeddings before Phase 1 encoding to avoid sync stall at Phase 2 start, improving timing accuracy and throughput
|
||||
- **🧹 MPS: Optimized model cleanup** - Skip redundant CPU movement before model deletion on unified memory
|
||||
|
||||
**2025.12.12 - Version 2.5.20**
|
||||
|
||||
- **⚡ Expanded attention backends** - Full support for Flash Attention 2 (Ampere+), Flash Attention 3 (Hopper+), SageAttention 2, and SageAttention 3 (Blackwell/RTX 50xx), with automatic fallback chains to PyTorch SDPA when unavailable *(based on PR by [@naxci1](https://github.com/naxci1) - thank you!)*
|
||||
@@ -805,14 +835,16 @@ python inference_cli.py image.jpg
|
||||
# Basic video upscaling with temporal consistency
|
||||
python inference_cli.py video.mp4 --resolution 720 --batch_size 33
|
||||
|
||||
# Streaming mode for long videos (memory-efficient)
|
||||
# Streaming mode for long videos (memory-efficient) with 10-bit video output (requires FFMPEG)
|
||||
# Processes video in chunks of 330 frames to avoid loading entire video into RAM
|
||||
# Use --temporal_overlap to ensure smooth transitions between chunks
|
||||
python inference_cli.py long_video.mp4 \
|
||||
--resolution 1080 \
|
||||
--batch_size 33 \
|
||||
--chunk_size 330 \
|
||||
--temporal_overlap 3
|
||||
--temporal_overlap 3 \
|
||||
--video_backend ffmpeg \
|
||||
--10bit
|
||||
|
||||
# Multi-GPU processing with temporal overlap
|
||||
python inference_cli.py video.mp4 \
|
||||
@@ -859,6 +891,8 @@ python inference_cli.py media_folder/ \
|
||||
- `<input>`: Input file (.mp4, .avi, .png, .jpg, etc.) or directory
|
||||
- `--output`: Output path (default: auto-generated in 'output/' directory)
|
||||
- `--output_format`: Output format: 'mp4' (video) or 'png' (image sequence). Default: auto-detect from input type
|
||||
- `--video_backend`: Video encoder backend: 'opencv' (default) or 'ffmpeg' (requires ffmpeg in PATH)
|
||||
- `--10bit`: Save 10-bit video with x265 codec and yuv420p10le pixel format (reduces banding in gradients). Without this flag, ffmpeg uses x264 (yuv420p) for maximum compatibility. Requires --video_backend ffmpeg
|
||||
- `--model_dir`: Model directory (default: ./models/SEEDVR2)
|
||||
|
||||
**Model Selection:**
|
||||
@@ -1012,7 +1046,7 @@ For detailed contribution guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md).
|
||||
|
||||
This ComfyUI implementation is a collaborative project by **[NumZ](https://github.com/numz)** and **[AInVFX](https://www.youtube.com/@AInVFX)** (Adrien Toupet), based on the original [SeedVR2](https://github.com/ByteDance-Seed/SeedVR) by ByteDance Seed Team.
|
||||
|
||||
Special thanks to our community contributors including [naxci1](https://github.com/naxci1), [benjaminherb](https://github.com/benjaminherb), [cmeka](https://github.com/cmeka), [FurkanGozukara](https://github.com/FurkanGozukara), [JohnAlcatraz](https://github.com/JohnAlcatraz), [lihaoyun6](https://github.com/lihaoyun6), [Luchuanzhao](https://github.com/Luchuanzhao), [Luke2642](https://github.com/Luke2642), [proxyid](https://github.com/proxyid), [q5sys](https://github.com/q5sys), and many others for their improvements, bug fixes, and testing.
|
||||
Special thanks to our community contributors including [naxci1](https://github.com/naxci1), [thehhmdb](https://github.com/thehhmdb), [s-cerevisiae](https://github.com/s-cerevisiae), [benjaminherb](https://github.com/benjaminherb), [cmeka](https://github.com/cmeka), [FurkanGozukara](https://github.com/FurkanGozukara), [JohnAlcatraz](https://github.com/JohnAlcatraz), [lihaoyun6](https://github.com/lihaoyun6), [Luchuanzhao](https://github.com/Luchuanzhao), [Luke2642](https://github.com/Luke2642), [proxyid](https://github.com/proxyid), [q5sys](https://github.com/q5sys), and many others for their improvements, bug fixes, and testing.
|
||||
|
||||
## 📜 License
|
||||
|
||||
|
||||
+110
-10
@@ -108,6 +108,8 @@ else:
|
||||
import torch
|
||||
import cv2
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import shutil
|
||||
|
||||
# Project imports
|
||||
from src.utils.downloads import download_weight
|
||||
@@ -118,7 +120,9 @@ from src.core.generation_utils import (
|
||||
prepare_runner,
|
||||
compute_generation_info,
|
||||
log_generation_start,
|
||||
blend_overlapping_frames
|
||||
blend_overlapping_frames,
|
||||
load_text_embeddings,
|
||||
script_directory
|
||||
)
|
||||
from src.core.generation_phases import (
|
||||
encode_all_batches,
|
||||
@@ -130,6 +134,81 @@ from src.utils.debug import Debug
|
||||
from src.optimization.memory_manager import clear_memory, get_gpu_backend, is_cuda_available
|
||||
debug = Debug(enabled=False) # Will be enabled via --debug CLI flag
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# FFMPEG Class
|
||||
# =============================================================================
|
||||
|
||||
class FFMPEGVideoWriter:
|
||||
"""
|
||||
Video writer using ffmpeg subprocess for encoding with 10-bit support.
|
||||
|
||||
Provides cv2.VideoWriter-compatible interface (write, isOpened, release) while
|
||||
using ffmpeg for encoding. Enables 10-bit output (yuv420p10le with x265) which
|
||||
reduces banding artifacts in gradients compared to 8-bit opencv output.
|
||||
|
||||
Args:
|
||||
path: Output video file path
|
||||
width: Frame width in pixels
|
||||
height: Frame height in pixels
|
||||
fps: Frames per second
|
||||
use_10bit: If True, uses x265 codec with yuv420p10le pixel format.
|
||||
If False, uses x264 with yuv420p (default: False)
|
||||
|
||||
Raises:
|
||||
RuntimeError: If ffmpeg is not found in system PATH
|
||||
|
||||
Note:
|
||||
Frames must be passed to write() in BGR format (same as cv2.VideoWriter).
|
||||
Internally converts to RGB for ffmpeg rawvideo input.
|
||||
"""
|
||||
|
||||
def __init__(self, path: str, width: int, height: int, fps: float, use_10bit: bool = False):
|
||||
pix_fmt = 'yuv420p10le' if use_10bit else 'yuv420p'
|
||||
codec = 'libx265' if use_10bit else 'libx264'
|
||||
|
||||
self.proc = subprocess.Popen(
|
||||
['ffmpeg', '-y', '-f', 'rawvideo', '-pix_fmt', 'rgb24',
|
||||
'-s', f'{width}x{height}', '-r', str(fps), '-i', '-',
|
||||
'-c:v', codec, '-pix_fmt', pix_fmt, '-preset', 'medium', '-crf', '12', path],
|
||||
stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
|
||||
)
|
||||
|
||||
def write(self, frame_bgr: np.ndarray):
|
||||
if not self.isOpened():
|
||||
raise RuntimeError("FFMPEGVideoWriter: ffmpeg process is not running")
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
try:
|
||||
self.proc.stdin.write(frame_rgb.astype(np.uint8).tobytes())
|
||||
self.proc.stdin.flush() # Critical: prevent buffering issues
|
||||
except BrokenPipeError:
|
||||
raise RuntimeError(
|
||||
"FFMPEGVideoWriter: ffmpeg process terminated unexpectedly. "
|
||||
"Check video path, codec support, and disk space."
|
||||
)
|
||||
|
||||
def isOpened(self) -> bool:
|
||||
return self.proc is not None and self.proc.poll() is None
|
||||
|
||||
def release(self):
|
||||
if self.proc:
|
||||
try:
|
||||
self.proc.stdin.close()
|
||||
except Exception:
|
||||
pass # Ignore errors on close
|
||||
|
||||
self.proc.wait()
|
||||
|
||||
if self.proc.returncode != 0:
|
||||
debug.log(
|
||||
f"ffmpeg exited with code {self.proc.returncode}. "
|
||||
"Check output file for corruption.",
|
||||
level="WARNING", force=True, category="file"
|
||||
)
|
||||
self.proc = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Device Management Helpers
|
||||
# =============================================================================
|
||||
@@ -445,7 +524,8 @@ def process_single_file(input_path: str, args: argparse.Namespace, device_list:
|
||||
if is_png:
|
||||
save_frames_to_image(result, output_path, base_name)
|
||||
else:
|
||||
video_writer = save_frames_to_video(result, output_path, fps)
|
||||
video_writer = save_frames_to_video(result, output_path, fps,
|
||||
video_backend=args.video_backend, use_10bit=args.use_10bit)
|
||||
if video_writer is not None:
|
||||
video_writer.release()
|
||||
|
||||
@@ -473,7 +553,8 @@ def process_single_file(input_path: str, args: argparse.Namespace, device_list:
|
||||
if is_png:
|
||||
save_frames_to_image(result, output_path, base_name, start_index=frames_written)
|
||||
else:
|
||||
video_writer = save_frames_to_video(result, output_path, fps, writer=video_writer)
|
||||
video_writer = save_frames_to_video(result, output_path, fps, writer=video_writer,
|
||||
video_backend=args.video_backend, use_10bit=args.use_10bit)
|
||||
|
||||
frames_written += result.shape[0]
|
||||
del result
|
||||
@@ -656,7 +737,9 @@ def save_frames_to_video(
|
||||
frames_tensor: torch.Tensor,
|
||||
output_path: str,
|
||||
fps: float = 30.0,
|
||||
writer: Optional[cv2.VideoWriter] = None
|
||||
writer: Optional[cv2.VideoWriter] = None,
|
||||
video_backend: str = "opencv",
|
||||
use_10bit: bool = False
|
||||
) -> Optional[cv2.VideoWriter]:
|
||||
"""
|
||||
Save frames tensor to MP4 video file.
|
||||
@@ -681,10 +764,13 @@ def save_frames_to_video(
|
||||
T, H, W, C = frames_np.shape
|
||||
|
||||
if writer is None:
|
||||
debug.log(f"Saving {T} frames to video: {output_path}", category="file")
|
||||
debug.log(f"Saving {T} frames to video: {output_path} (backend={video_backend})", category="file")
|
||||
os.makedirs(Path(output_path).parent, exist_ok=True)
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
writer = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
|
||||
if video_backend == "ffmpeg":
|
||||
writer = FFMPEGVideoWriter(output_path, W, H, fps, use_10bit)
|
||||
else:
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
writer = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
|
||||
if not writer.isOpened():
|
||||
raise ValueError(f"Cannot create video writer for: {output_path}")
|
||||
|
||||
@@ -858,6 +944,10 @@ def _process_frames_core(
|
||||
if runner_cache is not None:
|
||||
runner_cache['runner'] = runner
|
||||
|
||||
# Preload text embeddings before Phase 1 to avoid sync stall in Phase 2
|
||||
ctx['text_embeds'] = load_text_embeddings(script_directory, ctx['dit_device'], ctx['compute_dtype'], debug)
|
||||
debug.log("Loaded text embeddings for DiT", category="dit")
|
||||
|
||||
# Compute generation info and log start (handles prepending internally)
|
||||
frames_tensor, gen_info = compute_generation_info(
|
||||
ctx=ctx,
|
||||
@@ -1230,8 +1320,8 @@ Examples:
|
||||
Basic video upscaling with temporal consistency:
|
||||
python {invocation} video.mp4 --resolution 720 --batch_size 33
|
||||
|
||||
Streaming mode for long videos:
|
||||
python {invocation} long_video.mp4 --resolution 1080 --batch_size 33 --chunk_size 330 --temporal_overlap 3
|
||||
Streaming mode for long videos with 10-bit video output (requires FFMPEG):
|
||||
python {invocation} long_video.mp4 --resolution 1080 --batch_size 33 --chunk_size 330 --temporal_overlap 3 --video_backend ffmpeg --10bit
|
||||
|
||||
Multi-GPU processing with temporal overlap:
|
||||
python {invocation} video.mp4 --cuda_device 0,1 --resolution 1080 --batch_size 81 --uniform_batch_size --temporal_overlap 3 --prepend_frames 4
|
||||
@@ -1244,7 +1334,6 @@ Examples:
|
||||
|
||||
Batch directory processing:
|
||||
python {invocation} media_folder/ --output processed/ --cuda_device 0 --cache_dit --cache_vae --dit_offload_device cpu --vae_offload_device cpu --resolution 1080 --max_resolution 1920
|
||||
|
||||
"""
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
@@ -1262,6 +1351,11 @@ Examples:
|
||||
help="Output path (default: auto-generated in 'output/' directory)")
|
||||
io_group.add_argument("--output_format", type=str, default=None, choices=["mp4", "png", None],
|
||||
help="Output format: 'mp4' (video) or 'png' (image sequence). Default: auto-detect from input type")
|
||||
io_group.add_argument("--video_backend", type=str, default="opencv", choices=["opencv", "ffmpeg"],
|
||||
help="Video encoder backend: 'opencv' (default) or 'ffmpeg' (requires ffmpeg in PATH)")
|
||||
io_group.add_argument("--10bit", dest="use_10bit", action="store_true",
|
||||
help="Save 10-bit video with x265 codec (reduces banding). Without this flag, "
|
||||
"ffmpeg uses x264 for maximum compatibility. Requires --video_backend ffmpeg")
|
||||
io_group.add_argument("--model_dir", type=str, default=None,
|
||||
help=f"Model directory (default: ./models/{SEEDVR2_FOLDER_NAME})")
|
||||
|
||||
@@ -1438,6 +1532,12 @@ def main() -> None:
|
||||
debug.log(f"VAE decode tile overlap ({args.vae_decode_tile_overlap}) must be smaller than tile size ({args.vae_decode_tile_size})", level="ERROR", category="vae", force=True)
|
||||
sys.exit(1)
|
||||
|
||||
# Validate ffmpeg availability if selected
|
||||
if args.video_backend == "ffmpeg" and shutil.which("ffmpeg") is None:
|
||||
debug.log("--video_backend ffmpeg requires ffmpeg in PATH. Install ffmpeg or use --video_backend opencv",
|
||||
level="ERROR", category="setup", force=True)
|
||||
sys.exit(1)
|
||||
|
||||
# Inform about caching defaults
|
||||
if args.cache_dit and args.dit_offload_device == "none":
|
||||
offload_target = "system memory (CPU)" if get_gpu_backend() != "mps" else "unified memory"
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "seedvr2_videoupscaler"
|
||||
description = "SeedVR2 official ComfyUI integration: ByteDance-Seed's one-step diffusion-based video/image upscaling with memory-efficient inference"
|
||||
version = "2.5.20"
|
||||
version = "2.5.24"
|
||||
authors = [
|
||||
{name = "numz"},
|
||||
{name = "adrientoupet"}
|
||||
|
||||
@@ -337,13 +337,23 @@ def edge_guided_alpha_upscale(
|
||||
rgb_edges = detect_edges_batch(images=rgb_normalized, method='sobel', debug=debug)
|
||||
|
||||
# Step 1: Initial bicubic upscale provides smooth base before edge refinement
|
||||
alpha_upscaled = F.interpolate(
|
||||
input_alpha,
|
||||
size=(H_out, W_out),
|
||||
mode='bicubic',
|
||||
align_corners=False,
|
||||
antialias=True
|
||||
).clamp(0, 1)
|
||||
# MPS on PyTorch < 2.8 doesn't support bicubic+antialias - use CPU fallback
|
||||
try:
|
||||
alpha_upscaled = F.interpolate(
|
||||
input_alpha,
|
||||
size=(H_out, W_out),
|
||||
mode='bicubic',
|
||||
align_corners=False,
|
||||
antialias=True
|
||||
).clamp(0, 1)
|
||||
except NotImplementedError:
|
||||
alpha_upscaled = F.interpolate(
|
||||
input_alpha.cpu(),
|
||||
size=(H_out, W_out),
|
||||
mode='bicubic',
|
||||
align_corners=False,
|
||||
antialias=True
|
||||
).to(device).clamp(0, 1)
|
||||
|
||||
if is_binary_mask:
|
||||
if debug:
|
||||
|
||||
@@ -231,7 +231,11 @@ def encode_all_batches(
|
||||
if images is None:
|
||||
raise ValueError("Images to encode must be provided")
|
||||
else:
|
||||
ctx['input_images'] = images
|
||||
# MPS: keep on device to avoid sync overhead in Phase 4 color correction
|
||||
if ctx['vae_device'].type == 'mps' and images.device.type != 'mps':
|
||||
ctx['input_images'] = images.to(ctx['vae_device'])
|
||||
else:
|
||||
ctx['input_images'] = images
|
||||
|
||||
# Get total frame count from context (set in video_upscaler before encoding)
|
||||
total_frames = ctx.get('total_frames', len(images))
|
||||
@@ -860,7 +864,13 @@ def decode_all_batches(
|
||||
|
||||
# Pre-allocate final_video at the START of decode phase (before any batch processing)
|
||||
# This ensures we only need memory for final_video + 1 batch, not final_video + all batch_samples
|
||||
target_device = ctx['tensor_offload_device'] if ctx['tensor_offload_device'] is not None else 'cpu'
|
||||
# MPS: keep on device (unified memory, no benefit to CPU offload)
|
||||
if ctx['tensor_offload_device'] is not None:
|
||||
target_device = ctx['tensor_offload_device']
|
||||
elif ctx['vae_device'].type == 'mps':
|
||||
target_device = ctx['vae_device']
|
||||
else:
|
||||
target_device = 'cpu'
|
||||
channels_str = "RGBA" if C == 4 else "RGB"
|
||||
required_gb = (total_frames * true_h * true_w * C * 2) / (1024**3)
|
||||
debug.log(f"Pre-allocating output tensor: {total_frames} frames, {true_w}x{true_h}px, {channels_str} ({required_gb:.2f}GB)",
|
||||
|
||||
@@ -350,7 +350,12 @@ def setup_generation_context(
|
||||
vae_device = _normalize_device(vae_device)
|
||||
dit_offload_device = _normalize_device(dit_offload_device) if dit_offload_device is not None else None
|
||||
vae_offload_device = _normalize_device(vae_offload_device) if vae_offload_device is not None else None
|
||||
tensor_offload_device = _normalize_device(tensor_offload_device) if tensor_offload_device is not None else None
|
||||
# MPS unified memory: CPU offload causes sync overhead with no memory benefit
|
||||
is_mps = dit_device.type == 'mps' or vae_device.type == 'mps'
|
||||
if is_mps and tensor_offload_device is not None and str(tensor_offload_device) == 'cpu':
|
||||
tensor_offload_device = None
|
||||
else:
|
||||
tensor_offload_device = _normalize_device(tensor_offload_device) if tensor_offload_device is not None else None
|
||||
|
||||
# Set LOCAL_RANK to 0 for single-GPU inference mode
|
||||
# CLI multi-GPU uses CUDA_VISIBLE_DEVICES to restrict visibility per worker
|
||||
|
||||
@@ -146,7 +146,7 @@ def load_quantized_state_dict(checkpoint_path: str, device: torch.device = torch
|
||||
handle_prefix="model.diffusion_model."
|
||||
)
|
||||
elif checkpoint_path.endswith('.pth'):
|
||||
state = torch.load(checkpoint_path, map_location=device_str, mmap=True)
|
||||
state = torch.load(checkpoint_path, map_location=device_str, mmap=True, weights_only=True)
|
||||
else:
|
||||
raise ValueError(f"Unsupported checkpoint format. Expected .safetensors or .pth, got: {checkpoint_path}")
|
||||
|
||||
@@ -393,6 +393,20 @@ class GGUFTensor(torch.Tensor):
|
||||
if debug:
|
||||
debug.log(f"Error in {func.__name__} dequantization: {e}", level="WARNING", category="dit", force=True)
|
||||
raise
|
||||
|
||||
# Handle conv2d/conv3d operations (critical for GGUF VAE models)
|
||||
# Conv3d layers (InflatedCausalConv3d) are not replaced by layer replacement
|
||||
if func in {torch.nn.functional.conv2d, torch.nn.functional.conv3d}:
|
||||
if len(args) >= 2 and isinstance(args[1], cls): # weight is second arg
|
||||
try:
|
||||
weight_tensor = args[1]
|
||||
dequantized_weight = weight_tensor.dequantize(device=args[0].device, dtype=args[0].dtype)
|
||||
new_args = (args[0], dequantized_weight) + args[2:]
|
||||
return func(*new_args, **kwargs)
|
||||
except Exception as e:
|
||||
if debug:
|
||||
debug.log(f"Error in conv dequantization: {e}", level="WARNING", category="dit", force=True)
|
||||
raise
|
||||
|
||||
# For ALL other operations, delegate to parent WITHOUT dequantization
|
||||
# This includes .cpu(), .to(), .device, .dtype, .shape, etc.
|
||||
|
||||
@@ -19,7 +19,9 @@ from ..core.generation_utils import (
|
||||
setup_generation_context,
|
||||
prepare_runner,
|
||||
compute_generation_info,
|
||||
log_generation_start
|
||||
log_generation_start,
|
||||
load_text_embeddings,
|
||||
script_directory
|
||||
)
|
||||
from ..optimization.memory_manager import (
|
||||
cleanup_text_embeddings,
|
||||
@@ -437,6 +439,10 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
|
||||
# Store cache context in ctx for use in generation phases
|
||||
ctx['cache_context'] = cache_context
|
||||
|
||||
# Preload text embeddings before Phase 1 to avoid sync stall in Phase 2
|
||||
ctx['text_embeds'] = load_text_embeddings(script_directory, ctx['dit_device'], ctx['compute_dtype'], debug)
|
||||
debug.log("Loaded text embeddings for DiT", category="dit")
|
||||
|
||||
debug.log_memory_state("After model preparation", show_tensors=False, detailed_tensors=False)
|
||||
debug.end_timer("model_preparation", "Model preparation", force=True, show_breakdown=True)
|
||||
|
||||
|
||||
@@ -1093,7 +1093,7 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
|
||||
):
|
||||
extra_cond_dim = kwargs.pop("extra_cond_dim") if "extra_cond_dim" in kwargs else None
|
||||
self.slicing_sample_min_size = slicing_sample_min_size
|
||||
self.slicing_latent_min_size = slicing_sample_min_size // (2**temporal_scale_num)
|
||||
self.slicing_latent_min_size = max(1, slicing_sample_min_size // (2**temporal_scale_num))
|
||||
|
||||
super().__init__(
|
||||
in_channels=in_channels,
|
||||
@@ -1224,6 +1224,10 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
|
||||
|
||||
output = causal_conv_gather_outputs(output)
|
||||
|
||||
# MPS memory leak workaround (pytorch/pytorch#155060)
|
||||
if self.device.type == 'mps':
|
||||
torch.mps.empty_cache()
|
||||
|
||||
# Only transfer back if needed
|
||||
return output if output.device == x.device else output.to(x.device)
|
||||
|
||||
@@ -1240,6 +1244,10 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
|
||||
output = self.decoder(_z, memory_state=memory_state)
|
||||
output = causal_conv_gather_outputs(output)
|
||||
|
||||
# MPS memory leak workaround (pytorch/pytorch#155060)
|
||||
if self.device.type == 'mps':
|
||||
torch.mps.empty_cache()
|
||||
|
||||
# Only transfer back if needed
|
||||
return output if output.device == z.device else output.to(z.device)
|
||||
|
||||
@@ -1710,7 +1718,7 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
|
||||
if split_size is not None:
|
||||
self.enable_slicing()
|
||||
self.slicing_sample_min_size = split_size
|
||||
self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
|
||||
self.slicing_latent_min_size = max(1, split_size // self.temporal_downsample_factor)
|
||||
else:
|
||||
self.disable_slicing()
|
||||
for module in self.modules():
|
||||
|
||||
@@ -733,7 +733,7 @@ class VideoAutoencoderKL(nn.Module):
|
||||
if slicing_sample_min_size is None:
|
||||
slicing_sample_min_size = temporal_downsample_factor
|
||||
self.slicing_sample_min_size = slicing_sample_min_size
|
||||
self.slicing_latent_min_size = slicing_sample_min_size // (2**temporal_scale_num)
|
||||
self.slicing_latent_min_size = max(1, slicing_sample_min_size // (2**temporal_scale_num))
|
||||
|
||||
# pass init params to Encoder
|
||||
self.encoder = Encoder3D(
|
||||
@@ -886,7 +886,7 @@ class VideoAutoencoderKL(nn.Module):
|
||||
if split_size is not None:
|
||||
self.enable_slicing()
|
||||
self.slicing_sample_min_size = split_size
|
||||
self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
|
||||
self.slicing_latent_min_size = max(1, split_size // self.temporal_downsample_factor)
|
||||
else:
|
||||
self.disable_slicing()
|
||||
for module in self.modules():
|
||||
@@ -950,7 +950,7 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
|
||||
self.disable_slicing()
|
||||
self.slicing_sample_min_size = split_size
|
||||
if split_size is not None:
|
||||
self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
|
||||
self.slicing_latent_min_size = max(1, split_size // self.temporal_downsample_factor)
|
||||
for module in self.modules():
|
||||
if isinstance(module, InflatedCausalConv3d):
|
||||
module.set_memory_device(memory_device)
|
||||
@@ -98,8 +98,8 @@ def ensure_bitsandbytes_safe():
|
||||
try:
|
||||
import bitsandbytes
|
||||
# Success - bitsandbytes works, other nodes can use it
|
||||
except (ImportError, OSError, RuntimeError):
|
||||
# Installation broken or not present - create stub
|
||||
except (ImportError, OSError, RuntimeError, ValueError):
|
||||
# Installation broken, not present, or version detection failed - create stub
|
||||
stub = types.ModuleType('bitsandbytes')
|
||||
stub.__spec__ = importlib.machinery.ModuleSpec('bitsandbytes', None)
|
||||
stub.__file__ = None
|
||||
@@ -592,11 +592,11 @@ def validate_gguf_availability(operation: str = "load GGUF model", debug=None) -
|
||||
raise RuntimeError(f"GGUF library required to {operation}")
|
||||
|
||||
|
||||
# 4. NVIDIA Conv3d Memory Bug - Workaround for PyTorch 2.9-2.10 + cuDNN >= 91002
|
||||
# 4. NVIDIA Conv3d Memory Bug - Workaround for PyTorch >= 2.9 + cuDNN >= 91002
|
||||
def _check_conv3d_memory_bug():
|
||||
"""
|
||||
Check if Conv3d memory bug workaround needed.
|
||||
Bug: PyTorch 2.9-2.10 with cuDNN >= 91002 uses 3x memory for Conv3d
|
||||
Bug: PyTorch 2.9+ with cuDNN >= 91002 uses 3x memory for Conv3d
|
||||
with fp16/bfloat16 due to buggy dispatch layer.
|
||||
"""
|
||||
try:
|
||||
@@ -622,7 +622,8 @@ def _check_conv3d_memory_bug():
|
||||
parts = version_str.split('.')
|
||||
torch_version = tuple(int(p) for p in parts[:2])
|
||||
|
||||
if not ((2, 9) <= torch_version <= (2, 10)):
|
||||
# Bug affects PyTorch 2.9 and later versions
|
||||
if torch_version < (2, 9):
|
||||
return False
|
||||
|
||||
if not hasattr(torch.backends.cudnn, 'version'):
|
||||
@@ -826,8 +827,11 @@ class CompatibleDiT(torch.nn.Module):
|
||||
param.data = param.data.to(target_dtype)
|
||||
converted_count += 1
|
||||
|
||||
# Also convert buffers
|
||||
# Also convert buffers (skip GGUF quantized buffers - they have tensor_type attribute)
|
||||
for name, buffer in self.dit_model.named_buffers():
|
||||
# Skip GGUF quantized buffers - these must stay in packed format for on-the-fly dequantization
|
||||
if hasattr(buffer, 'tensor_type'):
|
||||
continue
|
||||
if buffer.dtype != target_dtype:
|
||||
if buffer.device.type == "mps":
|
||||
temp_cpu = buffer.data.to("cpu")
|
||||
|
||||
@@ -1050,15 +1050,17 @@ def cleanup_dit(runner: Any, debug: Optional['Debug'] = None, cache_model: bool
|
||||
|
||||
# Move model off GPU if needed
|
||||
if param_device.type not in ['meta', 'cpu']:
|
||||
# Get offload target - default to 'cpu' if not configured or set to 'none'
|
||||
offload_target = getattr(runner, '_dit_offload_device', None)
|
||||
if offload_target is None or offload_target == 'none':
|
||||
offload_target = torch.device('cpu')
|
||||
|
||||
# Move model off GPU (either for caching or before deletion)
|
||||
reason = "model caching" if cache_model else "releasing GPU memory"
|
||||
manage_model_device(model=runner.dit, target_device=offload_target, model_name="DiT",
|
||||
debug=debug, reason=reason, runner=runner)
|
||||
# MPS: skip CPU movement before deletion (unified memory, just causes sync)
|
||||
if param_device.type == 'mps' and not cache_model:
|
||||
if debug:
|
||||
debug.log("DiT on MPS - skipping CPU movement before deletion", category="cleanup")
|
||||
else:
|
||||
offload_target = getattr(runner, '_dit_offload_device', None)
|
||||
if offload_target is None or offload_target == 'none':
|
||||
offload_target = torch.device('cpu')
|
||||
reason = "model caching" if cache_model else "releasing GPU memory"
|
||||
manage_model_device(model=runner.dit, target_device=offload_target, model_name="DiT",
|
||||
debug=debug, reason=reason, runner=runner)
|
||||
elif param_device.type == 'meta' and debug:
|
||||
debug.log("DiT on meta device - keeping structure for cache", category="cleanup")
|
||||
except StopIteration:
|
||||
@@ -1126,15 +1128,17 @@ def cleanup_vae(runner: Any, debug: Optional['Debug'] = None, cache_model: bool
|
||||
|
||||
# Move model off GPU if needed
|
||||
if param_device.type not in ['meta', 'cpu']:
|
||||
# Get offload target - default to 'cpu' if not configured or set to 'none'
|
||||
offload_target = getattr(runner, '_vae_offload_device', None)
|
||||
if offload_target is None or offload_target == 'none':
|
||||
offload_target = torch.device('cpu')
|
||||
|
||||
# Move model off GPU (either for caching or before deletion)
|
||||
reason = "model caching" if cache_model else "releasing GPU memory"
|
||||
manage_model_device(model=runner.vae, target_device=offload_target, model_name="VAE",
|
||||
debug=debug, reason=reason, runner=runner)
|
||||
# MPS: skip CPU movement before deletion (unified memory, just causes sync)
|
||||
if param_device.type == 'mps' and not cache_model:
|
||||
if debug:
|
||||
debug.log("VAE on MPS - skipping CPU movement before deletion", category="cleanup")
|
||||
else:
|
||||
offload_target = getattr(runner, '_vae_offload_device', None)
|
||||
if offload_target is None or offload_target == 'none':
|
||||
offload_target = torch.device('cpu')
|
||||
reason = "model caching" if cache_model else "releasing GPU memory"
|
||||
manage_model_device(model=runner.vae, target_device=offload_target, model_name="VAE",
|
||||
debug=debug, reason=reason, runner=runner)
|
||||
elif param_device.type == 'meta' and debug:
|
||||
debug.log("VAE on meta device - keeping structure for cache", category="cleanup")
|
||||
except StopIteration:
|
||||
|
||||
+21
-9
@@ -381,6 +381,8 @@ def _rgb_to_lab_batch(rgb: Tensor, device: torch.device, matrix: Tensor, epsilon
|
||||
rgb_flat = rgb_linear.permute(0, 2, 3, 1).reshape(-1, 3)
|
||||
del rgb_linear
|
||||
|
||||
# Ensure dtype consistency for matrix multiplication
|
||||
rgb_flat = rgb_flat.to(dtype=matrix.dtype)
|
||||
xyz_flat = torch.matmul(rgb_flat, matrix.T)
|
||||
del rgb_flat
|
||||
|
||||
@@ -452,6 +454,8 @@ def _lab_to_rgb_batch(lab: Tensor, device: torch.device, matrix_inv: Tensor, eps
|
||||
xyz_flat = xyz.permute(0, 2, 3, 1).reshape(-1, 3)
|
||||
del xyz
|
||||
|
||||
# Ensure dtype consistency for matrix multiplication
|
||||
xyz_flat = xyz_flat.to(dtype=matrix_inv.dtype)
|
||||
rgb_linear_flat = torch.matmul(xyz_flat, matrix_inv.T)
|
||||
del xyz_flat
|
||||
|
||||
@@ -490,6 +494,7 @@ def _histogram_matching_channel(source: Tensor, reference: Tensor, device: torch
|
||||
# Sort both arrays
|
||||
source_sorted, source_indices = torch.sort(source_flat)
|
||||
reference_sorted, _ = torch.sort(reference_flat)
|
||||
del reference_flat
|
||||
|
||||
# Quantile mapping
|
||||
n_source = len(source_sorted)
|
||||
@@ -503,12 +508,15 @@ def _histogram_matching_channel(source: Tensor, reference: Tensor, device: torch
|
||||
ref_indices = (source_quantiles * (n_reference - 1)).long()
|
||||
ref_indices.clamp_(0, n_reference - 1)
|
||||
matched_sorted = reference_sorted[ref_indices]
|
||||
del source_quantiles, ref_indices
|
||||
del source_quantiles, ref_indices, reference_sorted
|
||||
|
||||
# Reconstruct with matched values
|
||||
matched_flat = torch.empty_like(source_flat)
|
||||
matched_flat.scatter_(0, source_indices, matched_sorted)
|
||||
del source_flat, reference_flat, source_sorted, source_indices, reference_sorted, matched_sorted
|
||||
del source_sorted, source_flat
|
||||
|
||||
# Reconstruct using argsort (portable across CUDA/ROCm/MPS)
|
||||
inverse_indices = torch.argsort(source_indices)
|
||||
del source_indices
|
||||
matched_flat = matched_sorted[inverse_indices]
|
||||
del matched_sorted, inverse_indices
|
||||
|
||||
return matched_flat.reshape(original_shape)
|
||||
|
||||
@@ -748,11 +756,15 @@ def _histogram_match_1d(source: Tensor, reference: Tensor, device: torch.device)
|
||||
ref_indices = (source_quantiles * (n_reference - 1)).long()
|
||||
ref_indices.clamp_(0, n_reference - 1)
|
||||
matched_sorted = reference_sorted[ref_indices]
|
||||
del source_quantiles, ref_indices
|
||||
del source_quantiles, ref_indices, reference_sorted
|
||||
|
||||
matched = torch.empty_like(source)
|
||||
matched.scatter_(0, source_indices, matched_sorted)
|
||||
del source_sorted, source_indices, reference_sorted, matched_sorted
|
||||
del source_sorted
|
||||
|
||||
# Reconstruct using argsort (portable across CUDA/ROCm/MPS)
|
||||
inverse_indices = torch.argsort(source_indices)
|
||||
del source_indices
|
||||
matched = matched_sorted[inverse_indices]
|
||||
del matched_sorted, inverse_indices
|
||||
|
||||
return matched
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ Only includes constants actually used in the codebase
|
||||
"""
|
||||
|
||||
# Version information
|
||||
__version__ = "2.5.20"
|
||||
__version__ = "2.5.24"
|
||||
|
||||
import os
|
||||
import warnings
|
||||
|
||||
+3
-3
@@ -78,7 +78,7 @@ class Debug:
|
||||
"device": "🖥️", # Device info
|
||||
"file": "📂", # File operations
|
||||
"alpha": "👻", # Alpha operations
|
||||
"star": "⭐", # Star
|
||||
"starlove": "⭐💝", # Star + love
|
||||
"dialogue": "💬", # Dialogue
|
||||
"none" : "",
|
||||
}
|
||||
@@ -259,9 +259,9 @@ class Debug:
|
||||
"""Print the footer with links - always displayed"""
|
||||
self.log("", category="none", force=True)
|
||||
self.log("────────────────────────", category="none", force=True)
|
||||
self.log("Questions? Updates? Watch the videos, star the repo & join us!", category="dialogue", force=True)
|
||||
self.log("Questions? Updates? Watch, star & sponsor if you can!", category="dialogue", force=True)
|
||||
self.log("https://www.youtube.com/@AInVFX", category="generation", force=True)
|
||||
self.log("https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler", category="star", force=True)
|
||||
self.log("https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler", category="starlove", force=True)
|
||||
|
||||
@torch._dynamo.disable # Skip tracing to avoid time.time() warnings
|
||||
def start_timer(self, name: str, force: bool = False) -> None:
|
||||
|
||||
Reference in New Issue
Block a user