Merge pull request #384 from AInVFX/main
v2.5.18: CLI streaming mode, multi-GPU streaming with caching, shared memory fix
This commit is contained in:
Adrien Toupet
@@ -36,6 +36,12 @@ We're actively working on improvements and new features. To stay informed:
|
||||
|
||||
## 🚀 Updates
|
||||
|
||||
**2025.12.09 - Version 2.5.18**
|
||||
|
||||
- **🚀 CLI: Streaming mode for long videos** - New `--chunk_size` flag processes videos in memory-bounded chunks, enabling arbitrarily long videos without RAM limits. Works with model caching (`--cache_dit`/`--cache_vae`) for chunk-to-chunk reuse *(inspired by [disk02](https://github.com/disk02) PR contribution)*
|
||||
- **⚡ CLI: Multi-GPU streaming** - Each GPU now streams its segment internally with independent model caching, improving memory efficiency and enabling `--temporal_overlap` blending at GPU boundaries
|
||||
- **🔧 CLI: Fix large video MemoryError** - Shared memory transfer replaces numpy pickling, preventing crashes on high-resolution/long video outputs *(inspired by [FurkanGozukara](https://github.com/FurkanGozukara) PR contribution)*
|
||||
|
||||
**2025.12.05 - Version 2.5.17**
|
||||
|
||||
- **🔧 Fix: Older GPU compatibility (GTX 970, etc.)** - Runtime bf16 CUBLAS probe replaces compute capability heuristics, correctly detecting unsupported GPUs without affecting RTX 20XX
|
||||
@@ -771,9 +777,18 @@ The CLI provides comprehensive options for single-GPU, multi-GPU, and batch proc
|
||||
# Basic image upscaling
|
||||
python inference_cli.py image.jpg
|
||||
|
||||
# Basic video video upscaling with temporal consistency
|
||||
# Basic video upscaling with temporal consistency
|
||||
python inference_cli.py video.mp4 --resolution 720 --batch_size 33
|
||||
|
||||
# Streaming mode for long videos (memory-efficient)
|
||||
# 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
|
||||
|
||||
# Multi-GPU processing with temporal overlap
|
||||
python inference_cli.py video.mp4 \
|
||||
--cuda_device 0,1 \
|
||||
@@ -830,7 +845,8 @@ python inference_cli.py media_folder/ \
|
||||
- `--batch_size`: Frames per batch (must follow 4n+1: 1, 5, 9, 13, 17, 21...). Ideally matches shot length for best temporal consistency (default: 5)
|
||||
- `--seed`: Random seed for reproducibility (default: 42)
|
||||
- `--skip_first_frames`: Skip N initial frames (default: 0)
|
||||
- `--load_cap`: Load maximum N frames from video. 0 = load all (default: 0)
|
||||
- `--load_cap`: Maximum total frames to load from video. 0 = load all (default: 0)
|
||||
- `--chunk_size`: Frames per chunk for streaming mode. When > 0, processes video in memory-bounded chunks of N frames, writing each chunk before loading the next. Essential for long videos that would otherwise exceed RAM. Use with `--temporal_overlap` for seamless chunk transitions. 0 = load all frames at once (default: 0)
|
||||
- `--prepend_frames`: Prepend N reversed frames to reduce start artifacts (auto-removed) (default: 0)
|
||||
- `--temporal_overlap`: Frames to overlap between batches/GPUs for smooth blending (default: 0)
|
||||
|
||||
|
||||
+489
-281
@@ -8,10 +8,12 @@ Supports single and multi-GPU processing with advanced memory optimization.
|
||||
Key Features:
|
||||
• Multi-GPU Processing: Automatic workload distribution across multiple GPUs with
|
||||
temporal overlap blending for seamless transitions
|
||||
• Streaming Mode: Memory-efficient processing of long videos in chunks, avoiding
|
||||
full video loading into RAM while maintaining temporal consistency
|
||||
• Memory Optimization: BlockSwap for limited VRAM, VAE tiling for large resolutions,
|
||||
intelligent tensor offloading between processing phases
|
||||
• Performance: Torch.compile integration, BFloat16 compute pipeline,
|
||||
efficient model caching for batch processing
|
||||
efficient model caching for batch and streaming processing
|
||||
• Flexibility: Multiple output formats (MP4/PNG), advanced color correction methods,
|
||||
directory batch processing with auto-format detection
|
||||
• Quality Control: Temporal overlap blending, frame prepending for artifact reduction,
|
||||
@@ -48,7 +50,7 @@ import argparse
|
||||
import time
|
||||
import platform
|
||||
import multiprocessing as mp
|
||||
from typing import Dict, Any, List, Optional, Tuple, Literal
|
||||
from typing import Dict, Any, List, Optional, Tuple, Literal, Generator
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
@@ -125,6 +127,7 @@ from src.core.generation_phases import (
|
||||
postprocess_all_batches
|
||||
)
|
||||
from src.utils.debug import Debug
|
||||
from src.optimization.memory_manager import clear_memory
|
||||
debug = Debug(enabled=False) # Will be enabled via --debug CLI flag
|
||||
|
||||
|
||||
@@ -356,6 +359,9 @@ def process_single_file(input_path: str, args: argparse.Namespace, device_list:
|
||||
"""
|
||||
Process a single video or image file with optional model caching.
|
||||
|
||||
For videos, supports streaming mode (chunk_size > 0) which processes in memory-bounded
|
||||
chunks with temporal overlap for seamless transitions between chunks.
|
||||
|
||||
Args:
|
||||
input_path: Path to input file
|
||||
args: Command-line arguments with all processing settings
|
||||
@@ -365,7 +371,7 @@ def process_single_file(input_path: str, args: argparse.Namespace, device_list:
|
||||
runner_cache: Optional cache dict for model reuse across multiple files
|
||||
|
||||
Returns:
|
||||
Number of frames processed from the input
|
||||
Number of frames written to output
|
||||
"""
|
||||
input_type = get_input_type(input_path)
|
||||
|
||||
@@ -375,19 +381,6 @@ def process_single_file(input_path: str, args: argparse.Namespace, device_list:
|
||||
|
||||
debug.log(f"Processing {input_type}: {Path(input_path).name}", category="generation", force=True)
|
||||
|
||||
# Extract frames
|
||||
if input_type == "video":
|
||||
start_time = time.time()
|
||||
frames_tensor, original_fps = extract_frames_from_video(
|
||||
input_path, args.skip_first_frames, args.load_cap
|
||||
)
|
||||
debug.log(f"Frame extraction time: {time.time() - start_time:.2f}s", category="timing")
|
||||
else:
|
||||
frames_tensor, original_fps = extract_frames_from_image(input_path)
|
||||
|
||||
# Track frames before processing (for FPS calculation)
|
||||
input_frame_count = len(frames_tensor)
|
||||
|
||||
# Generate or validate output path
|
||||
if output_path is None:
|
||||
output_path = generate_output_path(input_path, args.output_format, input_type=input_type)
|
||||
@@ -400,237 +393,360 @@ def process_single_file(input_path: str, args: argparse.Namespace, device_list:
|
||||
format_prefix = "Auto-detected" if format_auto_detected else "Requested"
|
||||
debug.log(f"{format_prefix} output format: {args.output_format}", category="info", force=True, indent_level=1)
|
||||
|
||||
# Process frames
|
||||
processing_start = time.time()
|
||||
# Use direct processing if caching enabled OR on Mac (MPS doesn't support multiprocessing well)
|
||||
if runner_cache is not None or platform.system() == "Darwin":
|
||||
# Direct single-GPU processing (required for Mac MPS, optional for caching)
|
||||
result = _single_gpu_direct_processing(frames_tensor, args, device_list[0], runner_cache)
|
||||
else:
|
||||
# Multi-GPU or non-cached processing via worker processes
|
||||
result = _gpu_processing(frames_tensor, device_list, args)
|
||||
debug.log(f"Processing time: {time.time() - processing_start:.2f}s", category="timing")
|
||||
|
||||
# Save results
|
||||
is_png_format = args.output_format == "png"
|
||||
is_single_image = input_type == "image"
|
||||
|
||||
if is_png_format and is_single_image:
|
||||
# Single PNG file
|
||||
os.makedirs(Path(output_path).parent, exist_ok=True)
|
||||
frame_np = (result[0].cpu().numpy() * 255.0).astype(np.uint8)
|
||||
# Convert RGB(A) to BGR(A) based on channel count
|
||||
if frame_np.shape[2] == 4:
|
||||
frame_save = cv2.cvtColor(frame_np, cv2.COLOR_RGBA2BGRA)
|
||||
# === VIDEO PROCESSING ===
|
||||
if input_type == "video":
|
||||
if not os.path.exists(input_path):
|
||||
raise FileNotFoundError(f"Video file not found: {input_path}")
|
||||
|
||||
cap = cv2.VideoCapture(input_path)
|
||||
if not cap.isOpened():
|
||||
raise ValueError(f"Cannot open video file: {input_path}")
|
||||
|
||||
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
|
||||
debug.log(f"Video info: {total_frames} frames, {width}x{height}, {fps:.2f} FPS", category="info")
|
||||
|
||||
# Skip initial frames
|
||||
if args.skip_first_frames > 0:
|
||||
debug.log(f"Skipping first {args.skip_first_frames} frames", category="info")
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, args.skip_first_frames)
|
||||
|
||||
# Calculate frames to process (apply load_cap if set)
|
||||
frames_to_process = total_frames - args.skip_first_frames
|
||||
if args.load_cap > 0:
|
||||
frames_to_process = min(frames_to_process, args.load_cap)
|
||||
|
||||
# Early exit for empty/exhausted video
|
||||
if frames_to_process <= 0:
|
||||
debug.log(f"No frames to process after skipping {args.skip_first_frames} of {total_frames}",
|
||||
level="WARNING", category="file", force=True)
|
||||
cap.release()
|
||||
return 0
|
||||
|
||||
# Streaming mode: process in chunks
|
||||
chunk_size = args.chunk_size if args.chunk_size > 0 else frames_to_process
|
||||
streaming = args.chunk_size > 0
|
||||
total_chunks = (frames_to_process + chunk_size - 1) // chunk_size # ceiling division
|
||||
|
||||
if streaming:
|
||||
debug.log(f"Streaming mode: chunks of {chunk_size} frames, overlap={args.temporal_overlap}",
|
||||
category="info", force=True, indent_level=1)
|
||||
|
||||
is_png = args.output_format == "png"
|
||||
video_writer = None
|
||||
overlap = args.temporal_overlap
|
||||
frames_written = 0
|
||||
chunk_idx = 0
|
||||
base_name = Path(input_path).stem
|
||||
|
||||
# Multi-GPU: workers stream their own segments
|
||||
if len(device_list) > 1:
|
||||
cap.release() # Workers will reopen
|
||||
video_info = {
|
||||
'video_path': input_path,
|
||||
'start_frame': args.skip_first_frames,
|
||||
'frames_to_process': frames_to_process,
|
||||
}
|
||||
result = _gpu_processing(None, device_list, args, video_info=video_info)
|
||||
|
||||
# Save result
|
||||
if is_png:
|
||||
save_frames_to_image(result, output_path, base_name)
|
||||
else:
|
||||
video_writer = save_frames_to_video(result, output_path, fps)
|
||||
if video_writer is not None:
|
||||
video_writer.release()
|
||||
|
||||
frames_written = result.shape[0]
|
||||
|
||||
# Single GPU: stream in main process
|
||||
else:
|
||||
frame_save = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR)
|
||||
cv2.imwrite(output_path, frame_save)
|
||||
|
||||
elif is_png_format:
|
||||
# PNG sequence (save_frames_to_png creates directory internally)
|
||||
save_frames_to_png(result, output_path, base_name=Path(input_path).stem)
|
||||
|
||||
else:
|
||||
# Video file
|
||||
os.makedirs(Path(output_path).parent, exist_ok=True)
|
||||
save_frames_to_video(result, output_path, original_fps)
|
||||
|
||||
# Log appropriate save message based on format
|
||||
if is_png_format and not is_single_image:
|
||||
debug.log(f"PNG frames saved in directory: {output_path}", category="file", force=True)
|
||||
else:
|
||||
chunk_count = 0
|
||||
for result in _stream_video_chunks(
|
||||
cap=cap,
|
||||
frames_to_process=frames_to_process,
|
||||
chunk_size=chunk_size,
|
||||
overlap=overlap,
|
||||
args=args,
|
||||
device_id=device_list[0],
|
||||
debug=debug,
|
||||
runner_cache=runner_cache,
|
||||
log_progress=streaming,
|
||||
total_chunks=total_chunks,
|
||||
cleanup_timer_name="chunk_cleanup"
|
||||
):
|
||||
chunk_count += 1
|
||||
|
||||
# Save output
|
||||
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)
|
||||
|
||||
frames_written += result.shape[0]
|
||||
del result
|
||||
|
||||
chunk_idx = chunk_count
|
||||
cap.release()
|
||||
if video_writer is not None:
|
||||
video_writer.release()
|
||||
|
||||
if streaming:
|
||||
debug.log("", category="none", force=True)
|
||||
if len(device_list) > 1:
|
||||
debug.log(f"Streaming complete: {frames_written} frames across {len(device_list)} GPUs", category="success", force=True)
|
||||
else:
|
||||
debug.log(f"Streaming complete: {frames_written} frames in {chunk_idx} chunks", category="success", force=True)
|
||||
|
||||
debug.log(f"Output saved to: {output_path}", category="file", force=True)
|
||||
return frames_written
|
||||
|
||||
return input_frame_count
|
||||
# === IMAGE PROCESSING ===
|
||||
frames_tensor, _ = extract_frames_from_image(input_path)
|
||||
|
||||
processing_start = time.time()
|
||||
# Process frames (multiprocessing only for multi-GPU)
|
||||
if len(device_list) > 1:
|
||||
result = _gpu_processing(frames_tensor, device_list, args)
|
||||
else:
|
||||
result = _single_gpu_direct_processing(frames_tensor, args, device_list[0], runner_cache)
|
||||
debug.log(f"Processing time: {time.time() - processing_start:.2f}s", category="timing")
|
||||
|
||||
# Save single image
|
||||
os.makedirs(Path(output_path).parent, exist_ok=True)
|
||||
frame_np = (result[0].cpu().numpy() * 255.0).astype(np.uint8)
|
||||
_save_image_bgr(frame_np, output_path)
|
||||
|
||||
debug.log(f"Output saved to: {output_path}", category="file", force=True)
|
||||
return 1
|
||||
|
||||
|
||||
def extract_frames_from_video(
|
||||
video_path: str,
|
||||
skip_first_frames: int = 0,
|
||||
load_cap: Optional[int] = None
|
||||
) -> Tuple[torch.Tensor, float]:
|
||||
def _read_frames_from_cap(cap: cv2.VideoCapture, max_frames: int) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Extract frames from video file and convert to tensor format.
|
||||
|
||||
Reads video using OpenCV, converts BGR to RGB, normalizes to [0,1] range.
|
||||
Note: Frame prepending is handled later in the processing pipeline via
|
||||
compute_generation_info(), not in this function.
|
||||
Read up to max_frames from an already-open VideoCapture.
|
||||
|
||||
Args:
|
||||
video_path: Path to input video file
|
||||
skip_first_frames: Number of initial frames to skip (default: 0)
|
||||
load_cap: Maximum number of frames to load, None loads all (default: None)
|
||||
|
||||
cap: An already opened cv2.VideoCapture instance
|
||||
max_frames: Maximum number of frames to read in this call
|
||||
|
||||
Returns:
|
||||
Tuple containing:
|
||||
- frames_tensor: Frames in format [T, H, W, C], Float32, range [0,1]
|
||||
- fps: Original video frames per second
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If video file doesn't exist
|
||||
ValueError: If video cannot be opened or no frames extracted
|
||||
Tensor [T, H, W, C] float32 [0,1], or None if no frames available
|
||||
"""
|
||||
debug.log(f"Extracting frames from video: {video_path}", category="file")
|
||||
|
||||
if not os.path.exists(video_path):
|
||||
raise FileNotFoundError(f"Video file not found: {video_path}")
|
||||
|
||||
# Open video
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
if not cap.isOpened():
|
||||
raise ValueError(f"Cannot open video file: {video_path}")
|
||||
|
||||
# Get video properties
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
|
||||
debug.log(f"Video info: {frame_count} frames, {width}x{height}, {fps:.2f} FPS", category="info")
|
||||
if skip_first_frames:
|
||||
debug.log(f"Will skip first {skip_first_frames} frames", category="info")
|
||||
if load_cap:
|
||||
debug.log(f"Will load maximum {load_cap} frames", category="info")
|
||||
|
||||
frames = []
|
||||
frame_idx = 0
|
||||
frames_loaded = 0
|
||||
|
||||
while True:
|
||||
for _ in range(max_frames):
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
# Skip first frame if requested
|
||||
if frame_idx < skip_first_frames:
|
||||
frame_idx += 1
|
||||
continue
|
||||
|
||||
if skip_first_frames > 0 and frame_idx == skip_first_frames:
|
||||
debug.log(f"Skipped first {skip_first_frames} frames", category="info")
|
||||
|
||||
# Check load cap
|
||||
if load_cap is not None and load_cap > 0 and frames_loaded >= load_cap:
|
||||
debug.log(f"Reached load cap of {load_cap} frames", category="info")
|
||||
break
|
||||
|
||||
# Convert BGR to RGB
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# Convert to float32 and normalize to 0-1
|
||||
frame = frame.astype(np.float32) / 255.0
|
||||
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
|
||||
frames.append(frame)
|
||||
frame_idx += 1
|
||||
frames_loaded += 1
|
||||
|
||||
if not frames:
|
||||
return None
|
||||
return torch.from_numpy(np.stack(frames)).to(torch.float32)
|
||||
|
||||
|
||||
def _stream_video_chunks(
|
||||
cap: cv2.VideoCapture,
|
||||
frames_to_process: int,
|
||||
chunk_size: int,
|
||||
overlap: int,
|
||||
args: argparse.Namespace,
|
||||
device_id: str,
|
||||
debug: 'Debug',
|
||||
runner_cache: Optional[Dict[str, Any]],
|
||||
log_progress: bool = False,
|
||||
total_chunks: int = 0,
|
||||
cleanup_timer_name: Optional[str] = None,
|
||||
log_prefix: str = ""
|
||||
) -> Generator[torch.Tensor, None, None]:
|
||||
"""
|
||||
Generator that streams and processes video chunks.
|
||||
|
||||
Handles frame reading, temporal context prepending, processing via
|
||||
_process_frames_core, context removal from output, and memory cleanup.
|
||||
Caller is responsible for VideoCapture lifecycle and result handling.
|
||||
|
||||
Args:
|
||||
cap: Open VideoCapture positioned at start frame
|
||||
frames_to_process: Total frames to read and process
|
||||
chunk_size: Frames per chunk (use frames_to_process for single chunk)
|
||||
overlap: Temporal overlap frames between chunks for blending
|
||||
args: Processing arguments (copied internally, prepend_frames zeroed after first chunk)
|
||||
device_id: GPU device ID for processing
|
||||
debug: Debug instance for logging
|
||||
runner_cache: Optional model cache dict for reuse across chunks
|
||||
log_progress: If True, log chunk progress with separators
|
||||
total_chunks: Total chunks for progress display (used if log_progress=True)
|
||||
cleanup_timer_name: Optional timer name for memory cleanup logging
|
||||
log_prefix: Optional prefix for log messages (e.g., "[GPU 0] " for worker identification)
|
||||
|
||||
Yields:
|
||||
Processed frames tensor [T, H, W, C] for each chunk, context frames removed
|
||||
"""
|
||||
chunk_args = argparse.Namespace(**vars(args))
|
||||
frames_read = 0
|
||||
prev_raw_tail = None
|
||||
chunk_idx = 0
|
||||
streaming = chunk_size < frames_to_process
|
||||
|
||||
while frames_read < frames_to_process:
|
||||
read_count = min(chunk_size, frames_to_process - frames_read)
|
||||
new_frames = _read_frames_from_cap(cap, read_count)
|
||||
if new_frames is None:
|
||||
break
|
||||
frames_read += new_frames.shape[0]
|
||||
chunk_idx += 1
|
||||
|
||||
if debug.enabled and frames_loaded % 100 == 0:
|
||||
total_to_load = min(frame_count, load_cap) if load_cap else frame_count
|
||||
debug.log(f"Extracted {frames_loaded}/{total_to_load} frames", category="file")
|
||||
|
||||
cap.release()
|
||||
|
||||
if len(frames) == 0:
|
||||
raise ValueError(f"No frames extracted from video: {video_path}")
|
||||
|
||||
debug.log(f"Extracted {len(frames)} frames", category="success")
|
||||
# Disable prepend_frames after first chunk
|
||||
if chunk_idx > 1:
|
||||
chunk_args.prepend_frames = 0
|
||||
|
||||
# Prepend context from previous chunk
|
||||
if prev_raw_tail is not None and overlap > 0:
|
||||
context_count = min(overlap, prev_raw_tail.shape[0])
|
||||
frames = torch.cat([prev_raw_tail[-context_count:], new_frames], dim=0)
|
||||
else:
|
||||
frames = new_frames
|
||||
context_count = 0
|
||||
|
||||
# Log progress if enabled
|
||||
if log_progress and streaming:
|
||||
if chunk_idx > 1:
|
||||
debug.log("", category="none", force=True)
|
||||
debug.log("━" * 60, category="none", force=True)
|
||||
debug.log("", category="none", force=True)
|
||||
debug.log(f"{log_prefix}Chunk {chunk_idx}/{total_chunks}: {new_frames.shape[0]} new + {context_count} context frames",
|
||||
category="generation", force=True)
|
||||
debug.log("", category="none", force=True)
|
||||
|
||||
# Process chunk
|
||||
result = _process_frames_core(
|
||||
frames_tensor=frames.to(torch.float16),
|
||||
args=chunk_args,
|
||||
device_id=device_id,
|
||||
debug=debug,
|
||||
runner_cache=runner_cache
|
||||
)
|
||||
|
||||
# Remove context frames from output
|
||||
if context_count > 0:
|
||||
result = result[context_count:]
|
||||
|
||||
# Save tail for next chunk context
|
||||
prev_raw_tail = new_frames[-overlap:].clone() if overlap > 0 else None
|
||||
|
||||
# Cleanup before yield
|
||||
del frames
|
||||
|
||||
yield result
|
||||
|
||||
# Memory cleanup between chunks
|
||||
if streaming:
|
||||
clear_memory(debug=debug, deep=True, force=True, timer_name=cleanup_timer_name)
|
||||
|
||||
# Convert to tensor (will be cast to compute_dtype in worker process)
|
||||
frames_tensor = torch.from_numpy(np.stack(frames)).to(torch.float32)
|
||||
|
||||
debug.log(f"Frames tensor shape: {frames_tensor.shape}, dtype: {frames_tensor.dtype}", category="memory")
|
||||
|
||||
return frames_tensor, fps
|
||||
def _save_image_bgr(frame_np: np.ndarray, file_path: str) -> None:
|
||||
"""
|
||||
Save a single RGB(A) uint8 frame to disk, converting to BGR(A) for OpenCV.
|
||||
|
||||
Args:
|
||||
frame_np: Frame as uint8 numpy array [H, W, C] where C is 3 (RGB) or 4 (RGBA)
|
||||
file_path: Output file path
|
||||
"""
|
||||
if frame_np.shape[2] == 4:
|
||||
frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGBA2BGRA)
|
||||
else:
|
||||
frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR)
|
||||
cv2.imwrite(file_path, frame_bgr)
|
||||
|
||||
|
||||
def save_frames_to_video(
|
||||
frames_tensor: torch.Tensor,
|
||||
output_path: str,
|
||||
fps: float = 30.0
|
||||
) -> None:
|
||||
fps: float = 30.0,
|
||||
writer: Optional[cv2.VideoWriter] = None
|
||||
) -> Optional[cv2.VideoWriter]:
|
||||
"""
|
||||
Save frames tensor to MP4 video file.
|
||||
|
||||
Converts tensor from Float32 [0,1] to uint8 [0,255], RGB to BGR for OpenCV,
|
||||
and writes to video file using mp4v codec.
|
||||
and writes to video file using mp4v codec. Supports streaming mode where
|
||||
an existing writer is passed and kept open for subsequent chunks.
|
||||
|
||||
Args:
|
||||
frames_tensor: Frames in format [T, H, W, C], Float32, range [0,1]
|
||||
output_path: Output video file path (will be created if doesn't exist)
|
||||
output_path: Output video file path (directory created if doesn't exist)
|
||||
fps: Frames per second for output video (default: 30.0)
|
||||
writer: Existing VideoWriter for streaming (if None, creates new one)
|
||||
|
||||
Returns:
|
||||
VideoWriter if streaming mode (caller must close), None if standalone mode
|
||||
|
||||
Raises:
|
||||
ValueError: If video writer cannot be initialized
|
||||
"""
|
||||
debug.log(f"Saving {frames_tensor.shape[0]} frames to video: {output_path}", category="file")
|
||||
|
||||
# Convert tensor to numpy and denormalize
|
||||
frames_np = frames_tensor.cpu().numpy()
|
||||
frames_np = (frames_np * 255.0).astype(np.uint8)
|
||||
|
||||
# Get video properties
|
||||
frames_np = (frames_tensor.cpu().numpy() * 255.0).astype(np.uint8)
|
||||
T, H, W, C = frames_np.shape
|
||||
|
||||
# Initialize video writer
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
|
||||
if writer is None:
|
||||
debug.log(f"Saving {T} frames to video: {output_path}", 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 not writer.isOpened():
|
||||
raise ValueError(f"Cannot create video writer for: {output_path}")
|
||||
|
||||
if not out.isOpened():
|
||||
raise ValueError(f"Cannot create video writer for: {output_path}")
|
||||
|
||||
# Write frames
|
||||
for i, frame in enumerate(frames_np):
|
||||
# Convert RGB to BGR for OpenCV
|
||||
frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
||||
out.write(frame_bgr)
|
||||
|
||||
writer.write(frame_bgr)
|
||||
if debug.enabled and (i + 1) % 100 == 0:
|
||||
debug.log(f"Saved {i + 1}/{T} frames", category="file")
|
||||
|
||||
out.release()
|
||||
debug.log(f"Written {i + 1}/{T} frames", category="file")
|
||||
|
||||
debug.log(f"Video saved successfully: {output_path}", category="success")
|
||||
return writer # Caller always closes
|
||||
|
||||
|
||||
def save_frames_to_png(
|
||||
def save_frames_to_image(
|
||||
frames_tensor: torch.Tensor,
|
||||
output_dir: str,
|
||||
base_name: str
|
||||
) -> None:
|
||||
base_name: str,
|
||||
start_index: int = 0
|
||||
) -> int:
|
||||
"""
|
||||
Save frames tensor as sequential PNG image files.
|
||||
|
||||
Each frame saved as {base_name}_{index:05d}.png with zero-padded indices.
|
||||
Each frame saved as {base_name}_{index:0Nd}.png with zero-padded indices.
|
||||
Converts Float32 [0,1] to uint8 [0,255] and RGB(A) to BGR(A) for OpenCV.
|
||||
|
||||
Args:
|
||||
frames_tensor: Frames in format [T, H, W, C], Float32, range [0,1]
|
||||
output_dir: Directory to save PNG files (created if doesn't exist)
|
||||
base_name: Base name for output files (e.g., "frame" → "frame_00000.png")
|
||||
start_index: Starting index for filenames (for streaming continuation)
|
||||
|
||||
Returns:
|
||||
Number of frames saved
|
||||
"""
|
||||
debug.log(f"Saving {frames_tensor.shape[0]} frames as PNGs to directory: {output_dir}", category="file")
|
||||
|
||||
# Ensure output directory exists
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Convert to numpy uint8 RGB
|
||||
|
||||
frames_np = (frames_tensor.cpu().numpy() * 255.0).astype(np.uint8)
|
||||
total = frames_np.shape[0]
|
||||
digits = max(5, len(str(total))) # at least 5 digits
|
||||
|
||||
if start_index == 0:
|
||||
debug.log(f"Saving {total} frames as PNGs to directory: {output_dir}", category="file")
|
||||
digits = 6 # Supports up to 999,999 frames (~11.5 hours at 24fps)
|
||||
|
||||
for idx, frame in enumerate(frames_np):
|
||||
filename = f"{base_name}_{idx:0{digits}d}.png"
|
||||
filename = f"{base_name}_{start_index + idx:0{digits}d}.png"
|
||||
file_path = os.path.join(output_dir, filename)
|
||||
# Convert RGB(A) to BGR(A) for cv2 based on channel count
|
||||
if frame.shape[2] == 4:
|
||||
frame_save = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGRA)
|
||||
else:
|
||||
frame_save = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
||||
cv2.imwrite(file_path, frame_save)
|
||||
_save_image_bgr(frame, file_path)
|
||||
if debug.enabled and (idx + 1) % 100 == 0:
|
||||
debug.log(f"Saved {idx + 1}/{total} PNGs", category="file")
|
||||
debug.log(f"Saved {idx + 1}/{total} images", category="file")
|
||||
|
||||
debug.log(f"PNG saving completed: {total} files in '{output_dir}'", category="success")
|
||||
debug.log(f"Saved {total} images to '{output_dir}'", category="success")
|
||||
return total
|
||||
|
||||
|
||||
# =============================================================================
|
||||
@@ -820,46 +936,100 @@ def _process_frames_core(
|
||||
def _worker_process(
|
||||
proc_idx: int,
|
||||
device_id: str,
|
||||
frames_np: np.ndarray,
|
||||
frames_np: Optional[np.ndarray],
|
||||
shared_args: Dict[str, Any],
|
||||
return_queue: mp.Queue
|
||||
return_queue: mp.Queue,
|
||||
done_barrier: mp.Barrier,
|
||||
video_info: Optional[Dict[str, Any]] = None
|
||||
) -> None:
|
||||
"""
|
||||
Worker process for multi-GPU upscaling.
|
||||
|
||||
CUDA_VISIBLE_DEVICES is set by parent before spawn, so this worker
|
||||
only sees its assigned GPU. Results returned via queue as numpy arrays.
|
||||
"""
|
||||
# Note: CUDA_VISIBLE_DEVICES and PYTORCH_CUDA_ALLOC_CONF are inherited
|
||||
# from parent (set before spawn). torch is imported at module level.
|
||||
Supports two modes:
|
||||
1. frames_np provided: Process pre-loaded frames (for images)
|
||||
2. video_info provided: Stream video segment internally (for videos)
|
||||
- Each worker opens the video, seeks to its assigned range, and streams
|
||||
with internal chunking and model caching for memory efficiency
|
||||
|
||||
Args:
|
||||
proc_idx: Worker index for result ordering
|
||||
device_id: GPU device ID (used for CUDA_VISIBLE_DEVICES inheritance)
|
||||
frames_np: Pre-loaded frames as numpy array, or None for video streaming
|
||||
shared_args: Serialized args namespace as dict
|
||||
return_queue: Queue for returning results to parent
|
||||
done_barrier: Barrier for synchronizing shared memory handoff
|
||||
video_info: Optional dict with 'video_path', 'start_frame', 'end_frame'
|
||||
for video streaming mode
|
||||
"""
|
||||
# Create debug instance for this worker
|
||||
worker_debug = Debug(enabled=shared_args["debug"])
|
||||
|
||||
# Convert numpy back to tensor
|
||||
frames_tensor = torch.from_numpy(frames_np).to(torch.float16)
|
||||
|
||||
# Create args namespace from shared_args
|
||||
args = argparse.Namespace(**shared_args)
|
||||
|
||||
# Process frames (no caching in worker mode)
|
||||
result_tensor = _process_frames_core(
|
||||
frames_tensor=frames_tensor,
|
||||
args=args,
|
||||
device_id="0", # Worker sees only 1 GPU (index 0) due to CUDA_VISIBLE_DEVICES
|
||||
debug=worker_debug,
|
||||
runner_cache=None # No caching in multiprocessing mode
|
||||
)
|
||||
# Video streaming mode: worker reads and processes its assigned segment
|
||||
if video_info is not None:
|
||||
cap = cv2.VideoCapture(video_info['video_path'])
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, video_info['start_frame'])
|
||||
|
||||
segment_frames = video_info['end_frame'] - video_info['start_frame']
|
||||
chunk_size = args.chunk_size if args.chunk_size > 0 else segment_frames
|
||||
|
||||
worker_debug.log(f"GPU {proc_idx}: frames {video_info['start_frame']}-{video_info['end_frame']} "
|
||||
f"({segment_frames} frames, chunks of {chunk_size})",
|
||||
category="generation", force=True)
|
||||
|
||||
# Only GPU 0 uses prepend_frames (applies to video start only)
|
||||
worker_args = argparse.Namespace(**vars(args))
|
||||
if proc_idx != 0:
|
||||
worker_args.prepend_frames = 0
|
||||
|
||||
# Enable model caching within worker only if requested
|
||||
runner_cache = {} if (args.cache_dit or args.cache_vae) else None
|
||||
|
||||
total_chunks = (segment_frames + chunk_size - 1) // chunk_size
|
||||
results = []
|
||||
for result in _stream_video_chunks(
|
||||
cap=cap,
|
||||
frames_to_process=segment_frames,
|
||||
chunk_size=chunk_size,
|
||||
overlap=args.temporal_overlap,
|
||||
args=worker_args,
|
||||
device_id="0",
|
||||
debug=worker_debug,
|
||||
runner_cache=runner_cache,
|
||||
log_progress=total_chunks > 1,
|
||||
total_chunks=total_chunks,
|
||||
log_prefix=f"[GPU {proc_idx}] "
|
||||
):
|
||||
results.append(result.cpu())
|
||||
|
||||
cap.release()
|
||||
result_tensor = torch.cat(results, dim=0) if results else torch.empty(0, dtype=torch.float32)
|
||||
|
||||
# Send back result as numpy array
|
||||
return_queue.put((proc_idx, result_tensor.numpy()))
|
||||
# Pre-loaded frames mode (original behavior)
|
||||
else:
|
||||
frames_tensor = torch.from_numpy(frames_np).to(torch.float16)
|
||||
result_tensor = _process_frames_core(
|
||||
frames_tensor=frames_tensor,
|
||||
args=args,
|
||||
device_id="0",
|
||||
debug=worker_debug,
|
||||
runner_cache=None
|
||||
)
|
||||
|
||||
# Share tensor memory for efficient cross-process transfer (avoids pickling large arrays)
|
||||
return_queue.put((proc_idx, result_tensor.share_memory_()))
|
||||
|
||||
# Wait for parent to copy shared tensors before exiting
|
||||
# (shared memory requires creating process to stay alive during access)
|
||||
done_barrier.wait()
|
||||
|
||||
|
||||
def _single_gpu_direct_processing(
|
||||
frames_tensor: torch.Tensor,
|
||||
args: argparse.Namespace,
|
||||
device_id: str,
|
||||
runner_cache: Dict[str, Any]
|
||||
runner_cache: Optional[Dict[str, Any]]
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Direct single-GPU processing with model caching support.
|
||||
@@ -876,88 +1046,115 @@ def _single_gpu_direct_processing(
|
||||
|
||||
|
||||
def _gpu_processing(
|
||||
frames_tensor: torch.Tensor,
|
||||
frames_tensor: Optional[torch.Tensor],
|
||||
device_list: List[str],
|
||||
args: argparse.Namespace
|
||||
args: argparse.Namespace,
|
||||
video_info: Optional[Dict[str, Any]] = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Orchestrate multi-GPU parallel video upscaling with temporal overlap blending.
|
||||
|
||||
Splits input frames across multiple GPUs with optional temporal overlap,
|
||||
spawns worker processes for parallel processing, and reassembles results
|
||||
with smooth blending of overlapping regions.
|
||||
|
||||
Processing flow:
|
||||
1. Split frames into chunks (with overlap if enabled)
|
||||
2. Spawn worker processes on each GPU
|
||||
3. Wait for all workers to complete
|
||||
4. Blend overlapping regions using Hann window crossfade
|
||||
5. Remove prepended frames from final result
|
||||
Supports two modes:
|
||||
1. video_info provided: Workers stream their assigned video segments internally
|
||||
(each GPU reads and processes its frame range with internal chunking)
|
||||
2. frames_tensor provided: Workers process pre-loaded frame chunks
|
||||
(non streaming behavior for images or pre-loaded videos)
|
||||
|
||||
Args:
|
||||
frames_tensor: Input frames [T, H, W, C], Float32, range [0,1]
|
||||
frames_tensor: Input frames [T, H, W, C] or None if using video_info mode
|
||||
device_list: List of device IDs as strings (e.g., ["0", "1"])
|
||||
args: Parsed command-line arguments containing all processing settings
|
||||
video_info: Optional dict with 'video_path', 'start_frame', 'frames_to_process'
|
||||
for streaming mode where workers read video directly
|
||||
|
||||
Returns:
|
||||
Upscaled frames tensor [T', H', W', C], Float32, range [0,1]
|
||||
where T' may be less than T if prepend_frames were removed
|
||||
|
||||
Note:
|
||||
- Single GPU: Can use multiprocessing or direct processing
|
||||
- Multi-GPU with overlap: Chunks sized to multiples of batch_size for
|
||||
proper temporal blending
|
||||
- Prepended frames removed after all GPU workers complete (multi-GPU safe)
|
||||
"""
|
||||
num_devices = len(device_list)
|
||||
total_frames = frames_tensor.shape[0]
|
||||
overlap = args.temporal_overlap
|
||||
|
||||
# Create overlapping chunks (for multi GPU); ensures every chunk is
|
||||
# a multiple of batch_size (except last one) to avoid blending issues
|
||||
if args.temporal_overlap > 0 and num_devices > 1:
|
||||
chunk_with_overlap = total_frames // num_devices + args.temporal_overlap
|
||||
if args.batch_size > 1:
|
||||
chunk_with_overlap = ((chunk_with_overlap + args.batch_size - 1) // args.batch_size) * args.batch_size
|
||||
base_chunk_size = chunk_with_overlap - args.temporal_overlap
|
||||
|
||||
chunks = []
|
||||
for i in range(num_devices):
|
||||
start_idx = i * base_chunk_size
|
||||
if i == num_devices - 1: # last chunk/device
|
||||
end_idx = total_frames
|
||||
else:
|
||||
end_idx = min(start_idx + chunk_with_overlap, total_frames)
|
||||
chunks.append(frames_tensor[start_idx:end_idx])
|
||||
else:
|
||||
chunks = torch.chunk(frames_tensor, num_devices, dim=0)
|
||||
|
||||
# Use direct Queue with explicit unlimited size for large video chunks
|
||||
return_queue = mp.Queue(maxsize=0) # 0 = unlimited (explicit)
|
||||
return_queue = mp.Queue(maxsize=0)
|
||||
done_barrier = mp.Barrier(num_devices + 1)
|
||||
workers = []
|
||||
|
||||
# Convert args namespace to dict for serialization
|
||||
shared_args = vars(args).copy()
|
||||
|
||||
# Start all workers
|
||||
for idx, (device_id, chunk_tensor) in enumerate(zip(device_list, chunks)):
|
||||
# Set CUDA_VISIBLE_DEVICES before spawning so child inherits it
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = device_id
|
||||
|
||||
# Video streaming mode: distribute frame ranges to workers
|
||||
if video_info is not None:
|
||||
total_frames = video_info['frames_to_process']
|
||||
start_frame = video_info['start_frame']
|
||||
video_path = video_info['video_path']
|
||||
|
||||
p = mp.Process(
|
||||
target=_worker_process,
|
||||
args=(idx, device_id, chunk_tensor.cpu().numpy(), shared_args, return_queue),
|
||||
)
|
||||
p.start()
|
||||
workers.append(p)
|
||||
base_per_gpu = total_frames // num_devices
|
||||
remainder = total_frames % num_devices
|
||||
|
||||
current_start = start_frame
|
||||
for idx, device_id in enumerate(device_list):
|
||||
gpu_frames = base_per_gpu + (1 if idx < remainder else 0)
|
||||
gpu_end = current_start + gpu_frames
|
||||
|
||||
# Add overlap frames for blending (except last GPU)
|
||||
if idx < num_devices - 1 and overlap > 0:
|
||||
gpu_end = min(gpu_end + overlap, start_frame + total_frames)
|
||||
|
||||
worker_video_info = {
|
||||
'video_path': video_path,
|
||||
'start_frame': current_start,
|
||||
'end_frame': gpu_end,
|
||||
}
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = device_id
|
||||
p = mp.Process(
|
||||
target=_worker_process,
|
||||
args=(idx, device_id, None, shared_args, return_queue, done_barrier),
|
||||
kwargs={'video_info': worker_video_info}
|
||||
)
|
||||
p.start()
|
||||
workers.append(p)
|
||||
|
||||
current_start += gpu_frames
|
||||
|
||||
# Pre-loaded frames mode (original behavior for images or non-streaming)
|
||||
else:
|
||||
total_frames = frames_tensor.shape[0]
|
||||
|
||||
if overlap > 0 and num_devices > 1:
|
||||
chunk_with_overlap = total_frames // num_devices + overlap
|
||||
if args.batch_size > 1:
|
||||
chunk_with_overlap = ((chunk_with_overlap + args.batch_size - 1) // args.batch_size) * args.batch_size
|
||||
base_chunk_size = chunk_with_overlap - overlap
|
||||
|
||||
chunks = []
|
||||
for i in range(num_devices):
|
||||
start_idx = i * base_chunk_size
|
||||
if i == num_devices - 1:
|
||||
end_idx = total_frames
|
||||
else:
|
||||
end_idx = min(start_idx + chunk_with_overlap, total_frames)
|
||||
chunks.append(frames_tensor[start_idx:end_idx])
|
||||
else:
|
||||
chunks = torch.chunk(frames_tensor, num_devices, dim=0)
|
||||
|
||||
for idx, (device_id, chunk_tensor) in enumerate(zip(device_list, chunks)):
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = device_id
|
||||
p = mp.Process(
|
||||
target=_worker_process,
|
||||
args=(idx, device_id, chunk_tensor.cpu().numpy(), shared_args, return_queue, done_barrier),
|
||||
)
|
||||
p.start()
|
||||
workers.append(p)
|
||||
|
||||
# Collect results before joining to prevent deadlock
|
||||
# Tensors arrive via shared memory - copy to numpy while workers still alive
|
||||
results_np = [None] * num_devices
|
||||
collected = 0
|
||||
while collected < num_devices:
|
||||
proc_idx, res_np = return_queue.get()
|
||||
results_np[proc_idx] = res_np
|
||||
proc_idx, result_tensor = return_queue.get()
|
||||
results_np[proc_idx] = result_tensor.numpy()
|
||||
collected += 1
|
||||
|
||||
# Release workers now that shared tensors are copied
|
||||
done_barrier.wait()
|
||||
|
||||
# Now safe to join
|
||||
for p in workers:
|
||||
p.join()
|
||||
@@ -1041,9 +1238,12 @@ Examples:
|
||||
Basic image upscaling:
|
||||
python {invocation} image.jpg
|
||||
|
||||
Basic video video upscaling with temporal consistency
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
@@ -1101,6 +1301,9 @@ Examples:
|
||||
help="Skip N initial frames (default: 0)")
|
||||
process_group.add_argument("--load_cap", type=int, default=0,
|
||||
help="Load maximum N frames from video. 0 = load all (default: 0)")
|
||||
process_group.add_argument("--chunk_size", type=int, default=0,
|
||||
help="Frames per chunk for streaming mode. When > 0, processes video in "
|
||||
"memory-bounded chunks of N frames. 0 = load all frames at once (default: 0)")
|
||||
process_group.add_argument("--prepend_frames", type=int, default=0,
|
||||
help="Prepend N reversed frames to reduce start artifacts (auto-removed). Default: 0")
|
||||
process_group.add_argument("--temporal_overlap", type=int, default=0,
|
||||
@@ -1325,18 +1528,18 @@ def main() -> None:
|
||||
|
||||
debug.log(f"Found {len(media_files)} media files to process", category="file", force=True)
|
||||
|
||||
# Validate caching with multi-GPU (not supported in CLI - would need shared memory)
|
||||
if (args.cache_dit or args.cache_vae) and len(device_list) > 1:
|
||||
# Multi-GPU caching requires streaming (workers cache within their chunk loops)
|
||||
if (args.cache_dit or args.cache_vae) and len(device_list) > 1 and args.chunk_size <= 0:
|
||||
debug.log(
|
||||
"Model caching requires single GPU selection (you selected multiple GPUs). "
|
||||
"Disabling caching for this run.",
|
||||
"Model caching requires streaming mode (--chunk_size > 0) for multi-GPU. "
|
||||
"Disabling caching for this run.",
|
||||
level="WARNING", category="cache", force=True
|
||||
)
|
||||
args.cache_dit = False
|
||||
args.cache_vae = False
|
||||
|
||||
# Initialize runner cache if caching enabled
|
||||
runner_cache = {} if (args.cache_dit or args.cache_vae) else None
|
||||
# Single-GPU: runner_cache persists across files; multi-GPU: workers cache internally
|
||||
runner_cache = {} if (args.cache_dit or args.cache_vae) and len(device_list) == 1 else None
|
||||
|
||||
for idx, file_path in enumerate(media_files, 1):
|
||||
# Visual separation between files (except before first file)
|
||||
@@ -1376,27 +1579,32 @@ def main() -> None:
|
||||
if format_auto_detected:
|
||||
args.output_format = "mp4" if input_type == "video" else "png"
|
||||
|
||||
# Validate caching for single file (would provide no benefit but shouldn't error)
|
||||
if (args.cache_dit or args.cache_vae):
|
||||
# Caching: single-GPU streaming uses runner_cache, multi-GPU streaming workers cache internally
|
||||
runner_cache = None
|
||||
streaming = args.chunk_size > 0
|
||||
|
||||
if args.cache_dit or args.cache_vae:
|
||||
if len(device_list) > 1:
|
||||
debug.log(
|
||||
"Model caching requires single GPU selection (you selected multiple GPUs). "
|
||||
"Disabling caching for this run.",
|
||||
level="WARNING", category="cache", force=True
|
||||
)
|
||||
args.cache_dit = False
|
||||
args.cache_vae = False
|
||||
if not streaming:
|
||||
debug.log(
|
||||
"Model caching requires streaming mode (--chunk_size > 0) for multi-GPU. "
|
||||
"Disabling caching for this run.",
|
||||
level="WARNING", category="cache", force=True
|
||||
)
|
||||
args.cache_dit = False
|
||||
args.cache_vae = False
|
||||
elif streaming:
|
||||
runner_cache = {}
|
||||
else:
|
||||
debug.log(
|
||||
"Model caching has no benefit for single file processing (only useful for directories). "
|
||||
"Model caching has no benefit for single file processing (only useful for directories or streaming mode). "
|
||||
"Consider removing --cache_dit/--cache_vae for single files.",
|
||||
category="tip", force=True
|
||||
)
|
||||
|
||||
# No caching for single file (no benefit)
|
||||
frames = process_single_file(args.input, args, device_list, args.output,
|
||||
format_auto_detected=format_auto_detected,
|
||||
runner_cache=None)
|
||||
runner_cache=runner_cache)
|
||||
total_frames_processed += frames
|
||||
|
||||
else:
|
||||
|
||||
+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.17"
|
||||
version = "2.5.18"
|
||||
authors = [
|
||||
{name = "numz"},
|
||||
{name = "adrientoupet"}
|
||||
|
||||
@@ -1414,7 +1414,7 @@ def postprocess_all_batches(
|
||||
total_computed += (num_valid_samples - 1) * actual_overlap
|
||||
frame_info += f" ({total_computed} computed with {' + '.join(adjustments)} removed)"
|
||||
|
||||
debug.log(f"Final output assembled: {frame_info}, Resolution: {Wf}x{Hf}px, Channels: {channels_str}",
|
||||
debug.log(f"Output assembled: {frame_info}, Resolution: {Wf}x{Hf}px, Channels: {channels_str}",
|
||||
category="generation", force=True)
|
||||
else:
|
||||
ctx['final_video'] = torch.empty((0, 0, 0, 0), dtype=ctx['compute_dtype'])
|
||||
|
||||
@@ -4,7 +4,7 @@ Only includes constants actually used in the codebase
|
||||
"""
|
||||
|
||||
# Version information
|
||||
__version__ = "2.5.17"
|
||||
__version__ = "2.5.18"
|
||||
|
||||
import os
|
||||
import warnings
|
||||
|
||||
Reference in New Issue
Block a user