diff --git a/README.md b/README.md index 6780481..07fd9b3 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,20 @@ We're actively working on improvements and new features. To stay informed: ## ๐Ÿš€ Release Notes +**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 +819,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 +875,8 @@ python inference_cli.py media_folder/ \ - ``: 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:** diff --git a/inference_cli.py b/inference_cli.py index 7c80200..24db7ea 100644 --- a/inference_cli.py +++ b/inference_cli.py @@ -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,63 @@ 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.PIPE + ) + + def write(self, frame_bgr: np.ndarray): + frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + self.proc.stdin.write(frame_rgb.astype(np.uint8).tobytes()) + + def isOpened(self) -> bool: + return self.proc is not None and self.proc.poll() is None + + def release(self): + if self.proc: + self.proc.stdin.close() + self.proc.wait() + stderr = self.proc.stderr.read() if self.proc.stderr else b'' + if self.proc.returncode != 0: + debug.log(f"ffmpeg error: {stderr.decode()}", level="WARNING", category="file") + self.proc = None + + # ============================================================================= # Device Management Helpers # ============================================================================= @@ -445,7 +506,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 +535,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 +719,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 +746,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 +926,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 +1302,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 +1316,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 +1333,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 +1514,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" diff --git a/pyproject.toml b/pyproject.toml index 61410dd..ff2ca79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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.22" authors = [ {name = "numz"}, {name = "adrientoupet"} diff --git a/src/core/alpha_upscaling.py b/src/core/alpha_upscaling.py index 67d3883..6059c7d 100644 --- a/src/core/alpha_upscaling.py +++ b/src/core/alpha_upscaling.py @@ -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: diff --git a/src/core/generation_phases.py b/src/core/generation_phases.py index 909b3cf..3b7e6ea 100644 --- a/src/core/generation_phases.py +++ b/src/core/generation_phases.py @@ -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)", diff --git a/src/core/generation_utils.py b/src/core/generation_utils.py index f89642e..9a4cb3c 100644 --- a/src/core/generation_utils.py +++ b/src/core/generation_utils.py @@ -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 diff --git a/src/interfaces/video_upscaler.py b/src/interfaces/video_upscaler.py index 54d6399..159ca2d 100644 --- a/src/interfaces/video_upscaler.py +++ b/src/interfaces/video_upscaler.py @@ -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) diff --git a/src/optimization/compatibility.py b/src/optimization/compatibility.py index cb386c8..09463c3 100644 --- a/src/optimization/compatibility.py +++ b/src/optimization/compatibility.py @@ -826,8 +826,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") diff --git a/src/optimization/memory_manager.py b/src/optimization/memory_manager.py index f5ca8af..780c690 100644 --- a/src/optimization/memory_manager.py +++ b/src/optimization/memory_manager.py @@ -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: diff --git a/src/utils/color_fix.py b/src/utils/color_fix.py index 7e9703b..880b9b9 100644 --- a/src/utils/color_fix.py +++ b/src/utils/color_fix.py @@ -490,6 +490,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 +504,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 +752,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 diff --git a/src/utils/constants.py b/src/utils/constants.py index d077e71..1b3882f 100644 --- a/src/utils/constants.py +++ b/src/utils/constants.py @@ -4,7 +4,7 @@ Only includes constants actually used in the codebase """ # Version information -__version__ = "2.5.20" +__version__ = "2.5.22" import os import warnings