Adrien Toupet
97c7b9cd12
Major improvements to device handling, model offloading, and code clarity: Device Management: - Replace string-based devices with torch.device objects throughout - Add explicit offload device parameters: dit_offload_device, vae_offload_device, tensor_offload_device - Remove preserve_vram in favor of explicit offload control - Improve get_device() to return torch.device objects consistently - Enhance get_device_list() with smart MPS-only system handling Context & Pipeline: - Merge setup_device_environment and prepare_generation_context into single setup_generation_context - Simplify LOCAL_RANK handling (set to '0' for single-GPU mode) - Store device configuration on runner for submodule access Tensor & Model Management: - Add manage_tensor_device() for consistent tensor movement with logging - Update manage_model_device() to use torch.device objects - Add validation for BlockSwap and caching configurations - Rename cache_in_ram to cache_model (more accurate naming) CLI & Interface: - Remove --preserve_vram flag (breaking change) - Add --vae_offload_device and --tensor_offload_device flags - Update ComfyUI node parameters with validation and better tooltips - Improve device selection UI with offload_device options
764 lines
35 KiB
Python
764 lines
35 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Standalone SeedVR2 Video Upscaler CLI Script
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import argparse
|
|
import time
|
|
import platform
|
|
import multiprocessing as mp
|
|
from typing import Dict, Any, List, Optional, Tuple
|
|
|
|
# Set up path before any other imports to fix module resolution
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
if script_dir not in sys.path:
|
|
sys.path.insert(0, script_dir)
|
|
|
|
# Set environment variable so all spawned processes can find modules
|
|
os.environ['PYTHONPATH'] = script_dir + ':' + os.environ.get('PYTHONPATH', '')
|
|
|
|
# Ensure safe CUDA usage with multiprocessing
|
|
if mp.get_start_method(allow_none=True) != 'spawn':
|
|
mp.set_start_method('spawn', force=True)
|
|
# -------------------------------------------------------------
|
|
# 1) Gestion VRAM (cudaMallocAsync) déjà en place
|
|
if platform.system() != "Darwin":
|
|
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
|
|
|
|
# 2) Pré-parse de la ligne de commande pour récupérer --cuda_device
|
|
_pre_parser = argparse.ArgumentParser(add_help=False)
|
|
_pre_parser.add_argument("--cuda_device", type=str, default=None)
|
|
_pre_args, _ = _pre_parser.parse_known_args()
|
|
if _pre_args.cuda_device is not None:
|
|
device_list_env = [x.strip() for x in _pre_args.cuda_device.split(',') if x.strip()!='']
|
|
if len(device_list_env) == 1:
|
|
# Single GPU: restrict visibility now
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = device_list_env[0]
|
|
|
|
# -------------------------------------------------------------
|
|
# 3) Imports lourds (torch, etc.) après la configuration env
|
|
import torch
|
|
import cv2
|
|
import numpy as np
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from src.utils.downloads import download_weight
|
|
from src.utils.debug import Debug
|
|
from src.utils.model_registry import get_available_dit_models, DEFAULT_DIT, DEFAULT_VAE
|
|
|
|
debug = Debug(enabled=False) # Default to disabled, can be enabled via CLI
|
|
|
|
|
|
def extract_frames_from_video(video_path: str, skip_first_frames: int = 0,
|
|
load_cap: Optional[int] = None, prepend_frames: int = 0) -> Tuple[torch.Tensor, float]:
|
|
"""
|
|
Extract frames from video and convert to tensor format
|
|
|
|
Args:
|
|
video_path (str): Path to input video
|
|
skip_first_frame (bool): Skip the first frame during extraction
|
|
load_cap (int): Maximum number of frames to load (None for all)
|
|
|
|
Returns:
|
|
torch.Tensor: Frames tensor in format [T, H, W, C] (Float16, normalized 0-1)
|
|
"""
|
|
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")
|
|
if prepend_frames:
|
|
debug.log(f"Will prepend {prepend_frames} frames to the video", category="info")
|
|
|
|
frames = []
|
|
frame_idx = 0
|
|
frames_loaded = 0
|
|
|
|
while True:
|
|
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
|
|
|
|
frames.append(frame)
|
|
frame_idx += 1
|
|
frames_loaded += 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")
|
|
|
|
# preprend frames if requested (reverse of the first few frames)
|
|
if prepend_frames > 0:
|
|
start_frames = []
|
|
if prepend_frames >= len(frames): # repeat first (=last) frame
|
|
start_frames = [frames[-1]] * (prepend_frames - len(frames) + 1)
|
|
frames = start_frames + frames[prepend_frames:0:-1] + frames
|
|
|
|
# Convert to tensor [T, H, W, C] and cast to Float16 for ComfyUI compatibility
|
|
frames_tensor = torch.from_numpy(np.stack(frames)).to(torch.float16)
|
|
|
|
debug.log(f"Frames tensor shape: {frames_tensor.shape}, dtype: {frames_tensor.dtype}", category="memory")
|
|
|
|
return frames_tensor, fps
|
|
|
|
|
|
def save_frames_to_video(frames_tensor: torch.Tensor, output_path: str, fps: float = 30.0) -> None:
|
|
"""
|
|
Save frames tensor to video file
|
|
|
|
Args:
|
|
frames_tensor (torch.Tensor): Frames in format [T, H, W, C] (Float16, 0-1)
|
|
output_path (str): Output video path
|
|
fps (float): Output video FPS
|
|
"""
|
|
debug.log(f"Saving {frames_tensor.shape[0]} frames to video: {output_path}", category="file")
|
|
|
|
# Ensure output directory exists
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
|
|
# Convert tensor to numpy and denormalize
|
|
frames_np = frames_tensor.cpu().numpy()
|
|
frames_np = (frames_np * 255.0).astype(np.uint8)
|
|
|
|
# Get video properties
|
|
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 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)
|
|
|
|
if debug.enabled and (i + 1) % 100 == 0:
|
|
debug.log(f"Saved {i + 1}/{T} frames", category="file")
|
|
|
|
out.release()
|
|
|
|
debug.log(f"Video saved successfully: {output_path}", category="success")
|
|
|
|
|
|
def save_frames_to_png(frames_tensor: torch.Tensor, output_dir: str, base_name: str) -> None:
|
|
"""
|
|
Save frames tensor as sequential PNG images.
|
|
|
|
Args:
|
|
frames_tensor (torch.Tensor): Frames in format [T, H, W, C] (Float16, 0-1)
|
|
output_dir (str): Directory to save PNGs
|
|
base_name (str): Base name for output files (without extension)
|
|
"""
|
|
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
|
|
|
|
for idx, frame in enumerate(frames_np):
|
|
filename = f"{base_name}_{idx:0{digits}d}.png"
|
|
file_path = os.path.join(output_dir, filename)
|
|
# Convert RGB to BGR for cv2
|
|
frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
|
cv2.imwrite(file_path, frame_bgr)
|
|
if debug.enabled and (idx + 1) % 100 == 0:
|
|
debug.log(f"Saved {idx + 1}/{total} PNGs", category="file")
|
|
|
|
debug.log(f"PNG saving completed: {total} files in '{output_dir}'", category="success")
|
|
|
|
|
|
def apply_temporal_overlap_blending(frames_tensor: torch.Tensor, batch_size: int, overlap: int) -> torch.Tensor:
|
|
"""
|
|
Blend frames with temporal overlap in pixel space and remove duplicates.
|
|
Args:
|
|
frames_tensor (torch.Tensor): [T, H, W, C], Float16 in [0,1]
|
|
batch_size (int): Frames per batch used during generation
|
|
overlap (int): Overlapping frames between consecutive batches
|
|
Returns:
|
|
torch.Tensor: Blended frames [T, H, W, C] with duplicates removed
|
|
"""
|
|
T = frames_tensor.shape[0]
|
|
if overlap <= 0 or batch_size <= overlap or T <= batch_size:
|
|
return frames_tensor
|
|
|
|
device = frames_tensor.device
|
|
dtype = frames_tensor.dtype
|
|
|
|
output = frames_tensor[:batch_size]
|
|
input_pos = batch_size
|
|
|
|
while input_pos < T:
|
|
remaining_frames = T - input_pos
|
|
current_batch_size = min(batch_size, remaining_frames)
|
|
|
|
if current_batch_size <= overlap:
|
|
break
|
|
|
|
current_batch = frames_tensor[input_pos:input_pos + current_batch_size]
|
|
|
|
prev_tail = output[-overlap:] # overlap frames from previous output
|
|
cur_head = current_batch[:overlap] # overlap frames from current batch
|
|
|
|
# Smooth crossfade while avoiding the first and last frames (which often have more artifacts)
|
|
if overlap >= 3:
|
|
t = torch.linspace(0.0, 1.0, steps=overlap, device=device, dtype=dtype)
|
|
blend_start = 1.0 / 3.0
|
|
blend_end = 2.0 / 3.0
|
|
u = ((t - blend_start) / (blend_end - blend_start)).clamp(0.0, 1.0)
|
|
w_prev_1d = 0.5 + 0.5 * torch.cos(torch.pi * u) # Hann window
|
|
else: # Linear fallback for small overlaps:
|
|
w_prev_1d = torch.linspace(1.0, 0, steps=overlap, device=device, dtype=dtype)
|
|
|
|
w_prev = w_prev_1d.view(overlap, 1, 1, 1)
|
|
w_cur = 1.0 - w_prev
|
|
blended = prev_tail * w_prev + cur_head * w_cur
|
|
|
|
# Replace the last overlap frames in output with blended result
|
|
output = torch.cat([output[:-overlap], blended], dim=0)
|
|
|
|
# Append the non-overlapping part of current batch (if any)
|
|
if overlap < current_batch_size:
|
|
non_overlapping = current_batch[overlap:]
|
|
output = torch.cat([output, non_overlapping], dim=0)
|
|
|
|
input_pos += current_batch_size
|
|
|
|
return output
|
|
|
|
|
|
def _worker_process(proc_idx: int, device_id: int, frames_np: np.ndarray,
|
|
shared_args: Dict[str, Any], return_queue: mp.Queue) -> None:
|
|
"""Worker process that performs upscaling on a slice of frames using a dedicated GPU."""
|
|
if platform.system() != "Darwin":
|
|
# 1. Limit CUDA visibility to the chosen GPU BEFORE importing torch-heavy deps
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
|
|
# Keep same cudaMallocAsync setting
|
|
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
|
|
|
|
import torch # local import inside subprocess
|
|
from src.core.generation import (
|
|
setup_generation_context, prepare_runner,
|
|
encode_all_batches, upscale_all_batches, decode_all_batches
|
|
)
|
|
|
|
# Create debug instance for this worker process
|
|
worker_debug = Debug(enabled=shared_args["debug"])
|
|
|
|
# Prepare offload device arguments
|
|
# DiT offload is disabled in CLI (no caching, single-run workflow)
|
|
vae_offload = None if shared_args["vae_offload_device"] == "none" else shared_args["vae_offload_device"]
|
|
tensor_offload = None if shared_args["tensor_offload_device"] == "none" else shared_args["tensor_offload_device"]
|
|
|
|
# Setup generation context with device configuration
|
|
ctx = setup_generation_context(
|
|
dit_device=f"cuda:{device_id}",
|
|
vae_device=f"cuda:{device_id}",
|
|
dit_offload_device=None,
|
|
vae_offload_device=vae_offload,
|
|
tensor_offload_device=tensor_offload,
|
|
debug=worker_debug
|
|
)
|
|
|
|
# Reconstruct frames tensor
|
|
frames_tensor = torch.from_numpy(frames_np).to(torch.float16)
|
|
|
|
# Create torch compile args if enabled
|
|
torch_compile_args_dit = None
|
|
torch_compile_args_vae = None
|
|
|
|
if shared_args.get("compile_dit", False):
|
|
torch_compile_args_dit = {
|
|
"backend": shared_args.get("compile_backend", "inductor"),
|
|
"mode": shared_args.get("compile_mode", "default"),
|
|
"fullgraph": shared_args.get("compile_fullgraph", False),
|
|
"dynamic": shared_args.get("compile_dynamic", False),
|
|
"dynamo_cache_size_limit": shared_args.get("compile_dynamo_cache_size_limit", 64),
|
|
"dynamo_recompile_limit": shared_args.get("compile_dynamo_recompile_limit", 128),
|
|
}
|
|
|
|
if shared_args.get("compile_vae", False):
|
|
torch_compile_args_vae = {
|
|
"backend": shared_args.get("compile_backend", "inductor"),
|
|
"mode": shared_args.get("compile_mode", "default"),
|
|
"fullgraph": shared_args.get("compile_fullgraph", False),
|
|
"dynamic": shared_args.get("compile_dynamic", False),
|
|
"dynamo_cache_size_limit": shared_args.get("compile_dynamo_cache_size_limit", 64),
|
|
"dynamo_recompile_limit": shared_args.get("compile_dynamo_recompile_limit", 128),
|
|
}
|
|
|
|
# Prepare runner
|
|
model_dir = shared_args["model_dir"]
|
|
model_name = shared_args["model"]
|
|
|
|
runner, _, _ = prepare_runner(
|
|
dit_model=model_name,
|
|
vae_model=DEFAULT_VAE,
|
|
model_dir=model_dir,
|
|
debug=worker_debug,
|
|
ctx=ctx,
|
|
dit_cache=False, # No caching in CLI
|
|
vae_cache=False, # No caching in CLI
|
|
dit_id=None, # No caching in CLI
|
|
vae_id=None, # No caching in CLI
|
|
block_swap_config=shared_args["block_swap_config"],
|
|
encode_tiled=shared_args["vae_encode_tiling_enabled"],
|
|
encode_tile_size=shared_args["vae_encode_tile_size"],
|
|
encode_tile_overlap=shared_args["vae_encode_tile_overlap"],
|
|
decode_tiled=shared_args["vae_decode_tiling_enabled"],
|
|
decode_tile_size=shared_args["vae_decode_tile_size"],
|
|
decode_tile_overlap=shared_args["vae_decode_tile_overlap"],
|
|
torch_compile_args_dit=torch_compile_args_dit,
|
|
torch_compile_args_vae=torch_compile_args_vae
|
|
)
|
|
|
|
# Phase 1: Encode all batches
|
|
ctx = encode_all_batches(
|
|
runner,
|
|
ctx=ctx,
|
|
images=frames_tensor,
|
|
batch_size=shared_args["batch_size"],
|
|
debug=worker_debug,
|
|
progress_callback=None,
|
|
temporal_overlap=shared_args["temporal_overlap"],
|
|
res_w=shared_args["res_w"],
|
|
input_noise_scale=shared_args["input_noise_scale"],
|
|
color_correction=shared_args.get("color_correction", "wavelet")
|
|
)
|
|
|
|
# Phase 2: Upscale all batches
|
|
ctx = upscale_all_batches(
|
|
runner,
|
|
ctx=ctx,
|
|
debug=worker_debug,
|
|
progress_callback=None,
|
|
cfg_scale=shared_args["cfg_scale"],
|
|
seed=shared_args["seed"],
|
|
latent_noise_scale=shared_args["latent_noise_scale"],
|
|
cache_model=False # No caching in CLI
|
|
)
|
|
|
|
# Phase 3: Decode all batches
|
|
ctx = decode_all_batches(
|
|
runner,
|
|
ctx=ctx,
|
|
debug=worker_debug,
|
|
progress_callback=None,
|
|
cache_model=False # No caching in CLI
|
|
)
|
|
|
|
# Phase 4: Post-processing and final assembly
|
|
ctx = postprocess_all_batches(
|
|
ctx=ctx,
|
|
debug=worker_debug,
|
|
progress_callback=None,
|
|
color_correction=shared_args.get("color_correction", "wavelet")
|
|
)
|
|
|
|
# Get final result
|
|
result_tensor = ctx['final_video']
|
|
|
|
# Ensure result is on CPU before converting to numpy
|
|
if result_tensor.is_cuda or result_tensor.is_mps:
|
|
result_tensor = result_tensor.cpu()
|
|
|
|
# Send back result as numpy array to avoid CUDA transfers
|
|
return_queue.put((proc_idx, result_tensor.numpy()))
|
|
|
|
|
|
def _gpu_processing(frames_tensor: torch.Tensor, device_list: List[str],
|
|
args: argparse.Namespace) -> torch.Tensor:
|
|
"""Split frames and process them in parallel on multiple GPUs."""
|
|
num_devices = len(device_list)
|
|
total_frames = frames_tensor.shape[0]
|
|
|
|
# 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)
|
|
|
|
manager = mp.Manager()
|
|
return_queue = manager.Queue()
|
|
workers = []
|
|
|
|
shared_args = {
|
|
"model": args.model,
|
|
"model_dir": args.model_dir if args.model_dir is not None else "./models/SEEDVR2",
|
|
"color_correction": args.color_correction,
|
|
"input_noise_scale": args.input_noise_scale,
|
|
"latent_noise_scale": args.latent_noise_scale,
|
|
"debug": args.debug,
|
|
"cfg_scale": 1.0,
|
|
"seed": args.seed,
|
|
"res_w": args.resolution,
|
|
"batch_size": args.batch_size,
|
|
"temporal_overlap": args.temporal_overlap,
|
|
"block_swap_config": {
|
|
'blocks_to_swap': args.blocks_to_swap,
|
|
'swap_io_components': args.swap_io_components,
|
|
},
|
|
"vae_encode_tiling_enabled": args.vae_encode_tiling_enabled,
|
|
"vae_encode_tile_size": args.vae_encode_tile_size,
|
|
"vae_encode_tile_overlap": args.vae_encode_tile_overlap,
|
|
"vae_decode_tiling_enabled": args.vae_decode_tiling_enabled,
|
|
"vae_decode_tile_size": args.vae_decode_tile_size,
|
|
"vae_decode_tile_overlap": args.vae_decode_tile_overlap,
|
|
"vae_offload_device": args.vae_offload_device,
|
|
"tensor_offload_device": args.tensor_offload_device,
|
|
"compile_dit": args.compile_dit,
|
|
"compile_vae": args.compile_vae,
|
|
"compile_backend": args.compile_backend,
|
|
"compile_mode": args.compile_mode,
|
|
"compile_fullgraph": args.compile_fullgraph,
|
|
"compile_dynamic": args.compile_dynamic,
|
|
"compile_dynamo_cache_size_limit": args.compile_dynamo_cache_size_limit,
|
|
"compile_dynamo_recompile_limit": args.compile_dynamo_recompile_limit,
|
|
}
|
|
|
|
for idx, (device_id, chunk_tensor) in enumerate(zip(device_list, chunks)):
|
|
p = mp.Process(
|
|
target=_worker_process,
|
|
args=(idx, device_id, chunk_tensor.cpu().numpy(), shared_args, return_queue),
|
|
)
|
|
p.start()
|
|
workers.append(p)
|
|
|
|
results_np = [None] * num_devices
|
|
collected = 0
|
|
while collected < num_devices:
|
|
proc_idx, res_np = return_queue.get()
|
|
results_np[proc_idx] = res_np
|
|
collected += 1
|
|
|
|
for p in workers:
|
|
p.join()
|
|
|
|
# Concatenate results with overlap handling
|
|
if args.temporal_overlap > 0 and num_devices > 1:
|
|
# Reconstruct results considering overlap
|
|
result_list = []
|
|
overlap = args.temporal_overlap
|
|
|
|
for idx, res_np in enumerate(results_np):
|
|
if idx == 0:
|
|
# First chunk: keep all frames
|
|
result_list.append(torch.from_numpy(res_np).to(torch.float16))
|
|
elif idx == num_devices - 1:
|
|
# Last chunk: skip overlap frames at the beginning
|
|
chunk_tensor = torch.from_numpy(res_np).to(torch.float16)
|
|
if chunk_tensor.shape[0] > overlap:
|
|
result_list.append(chunk_tensor[overlap:])
|
|
else:
|
|
# If chunk is smaller than overlap, skip it entirely
|
|
pass
|
|
else:
|
|
# Middle chunks: skip overlap at beginning, keep overlap at end
|
|
chunk_tensor = torch.from_numpy(res_np).to(torch.float16)
|
|
if chunk_tensor.shape[0] > overlap:
|
|
result_list.append(chunk_tensor[overlap:])
|
|
|
|
if result_list:
|
|
result_tensor = torch.cat(result_list, dim=0)
|
|
else:
|
|
result_tensor = torch.from_numpy(results_np[0]).to(torch.float16)
|
|
else:
|
|
# Original concatenation without overlap handling
|
|
result_tensor = torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float16)
|
|
|
|
return result_tensor
|
|
|
|
|
|
class OneOrTwoValues(argparse.Action):
|
|
def __call__(self, parser, namespace, values, option_string=None):
|
|
if len(values) not in [1, 2]:
|
|
parser.error(f"{option_string} requires 1 or 2 arguments")
|
|
|
|
if len(values) == 1:
|
|
values = values[0]
|
|
if ',' in values:
|
|
values = [v.strip() for v in values.split(',') if v.strip()]
|
|
else:
|
|
values = values.split()
|
|
|
|
try:
|
|
result = tuple(int(v) for v in values)
|
|
if len(result) == 1:
|
|
result = (result[0], result[0]) # Convert single value to (h, w)
|
|
setattr(namespace, self.dest, result)
|
|
except ValueError:
|
|
parser.error(f"{option_string} arguments must be integers")
|
|
|
|
|
|
def parse_arguments() -> argparse.Namespace:
|
|
"""Parse command line arguments"""
|
|
parser = argparse.ArgumentParser(description="SeedVR2 Video Upscaler CLI")
|
|
|
|
parser.add_argument("--video_path", type=str, required=True,
|
|
help="Path to input video file")
|
|
parser.add_argument("--seed", type=int, default=42,
|
|
help="Random seed for generation (default: 42)")
|
|
parser.add_argument("--resolution", type=int, default=1072,
|
|
help="Target resolution of the short side (default: 1072)")
|
|
parser.add_argument("--batch_size", type=int, default=1,
|
|
help="Number of frames per batch (default: 1)")
|
|
parser.add_argument("--model", type=str, default=DEFAULT_DIT,
|
|
choices=get_available_dit_models(),
|
|
help="Model to use (default: 3B FP8)")
|
|
parser.add_argument("--model_dir", type=str, default="seedvr2_models",
|
|
help="Directory containing the model files (default: use cache directory)")
|
|
parser.add_argument("--skip_first_frames", type=int, default=0,
|
|
help="Skip the first frames during processing")
|
|
parser.add_argument("--load_cap", type=int, default=0,
|
|
help="Maximum number of frames to load from video (default: load all)")
|
|
parser.add_argument("--output", type=str, default=None,
|
|
help="Output path (default: auto-generated, if output_format is png, it will be a directory)")
|
|
parser.add_argument("--output_format", type=str, default="video", choices=["video", "png"],
|
|
help="Output format: 'video' (mp4) or 'png' images (default: video)")
|
|
parser.add_argument("--color_correction", type=str, default="lab",
|
|
choices=["lab", "wavelet", "wavelet_adaptive", "hsv", "adain", "none"],
|
|
help="Color correction method: 'lab' (full perceptual color matching with detail preservation, recommended), 'wavelet' (frequency-based natural colors, preserves details), 'wavelet_adaptive' (wavelet base + targeted saturation correction), 'hsv' (hue-conditional saturation matching), 'adain' (statistical style transfer), 'none' (no correction)")
|
|
parser.add_argument("--input_noise_scale", type=float, default=0.0,
|
|
help="Input noise scale (0.0-1.0) to reduce artifacts at high resolutions. (default: 0.0)")
|
|
parser.add_argument("--latent_noise_scale", type=float, default=0.0,
|
|
help="Latent space noise scale (0.0-1.0). Adds noise during diffusion, can soften details. Use if input_noise doesn't help (default: 0.0)")
|
|
parser.add_argument("--debug", action="store_true",
|
|
help="Enable debug logging")
|
|
if platform.system() != "Darwin":
|
|
parser.add_argument("--cuda_device", type=str, default=None,
|
|
help="CUDA device id(s). Single id (e.g., '0') or comma-separated list '0,1' for multi-GPU")
|
|
parser.add_argument("--blocks_to_swap", type=int, default=0,
|
|
help="Number of blocks to swap for VRAM optimization (default: 0, disabled), up to 32 for 3B model, 36 for 7B")
|
|
parser.add_argument("--use_none_blocking", action="store_true",
|
|
help="Use non-blocking memory transfers for VRAM optimization")
|
|
parser.add_argument("--temporal_overlap", type=int, default=0,
|
|
help="Temporal overlap for processing (default: 0, no temporal overlap)")
|
|
parser.add_argument("--prepend_frames", type=int, default=0,
|
|
help="Number of frames to prepend to the video (default: 0). This can help with artifacts at the start of the video and are removed after processing")
|
|
parser.add_argument("--swap_io_components", action="store_true",
|
|
help="Swap IO components to CPU for VRAM optimization")
|
|
parser.add_argument("--vae_offload_device", type=str, default="none",
|
|
help="Device to offload VAE when not in use (default: none). "
|
|
"Options: 'cpu', 'none'. Use 'cpu' to free VRAM between encode/decode phases (slower but saves VRAM), "
|
|
"'none' to keep VAE on GPU throughout (faster but uses more VRAM)")
|
|
parser.add_argument("--tensor_offload_device", type=str, default="cpu",
|
|
help="Device to offload intermediate tensors between phases (default: cpu). "
|
|
"Options: 'cpu', 'none'. Use 'cpu' to prevent VRAM accumulation for long videos (recommended), "
|
|
"'none' to keep all tensors on GPU (faster but uses more VRAM)")
|
|
parser.add_argument("--vae_encode_tiling_enabled", action="store_true",
|
|
help="Enable VAE encode tiling for improved VRAM usage")
|
|
parser.add_argument("--vae_encode_tile_size", action=OneOrTwoValues, nargs='+', default=(512, 512),
|
|
help="VAE encode tile size (default: 512). Use single integer or two integers 'h w'. Only used if --vae_encode_tiling_enabled is set")
|
|
parser.add_argument("--vae_encode_tile_overlap", action=OneOrTwoValues, nargs='+', default=(128, 128),
|
|
help="VAE encode tile overlap (default: 128). Use single integer or two integers 'h w'. Only used if --vae_encode_tiling_enabled is set")
|
|
parser.add_argument("--vae_decode_tiling_enabled", action="store_true",
|
|
help="Enable VAE decode tiling for improved VRAM usage")
|
|
parser.add_argument("--vae_decode_tile_size", action=OneOrTwoValues, nargs='+', default=(512, 512),
|
|
help="VAE decode tile size (default: 512). Use single integer or two integers 'h w'. Only used if --vae_decode_tiling_enabled is set")
|
|
parser.add_argument("--vae_decode_tile_overlap", action=OneOrTwoValues, nargs='+', default=(128, 128),
|
|
help="VAE decode tile overlap (default: 128). Use single integer or two integers 'h w'. Only used if --vae_decode_tiling_enabled is set")
|
|
parser.add_argument("--compile_dit", action="store_true",
|
|
help="Enable torch.compile for DiT model (20-40%% speedup, requires PyTorch 2.0+)")
|
|
parser.add_argument("--compile_vae", action="store_true",
|
|
help="Enable torch.compile for VAE model (15-25%% speedup, requires PyTorch 2.0+)")
|
|
parser.add_argument("--compile_backend", type=str, default="inductor", choices=["inductor", "cudagraphs"],
|
|
help="Torch compile backend (default: inductor)")
|
|
parser.add_argument("--compile_mode", type=str, default="default", choices=["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"],
|
|
help="Torch compile mode (default: default)")
|
|
parser.add_argument("--compile_fullgraph", action="store_true",
|
|
help="Compile entire model as single graph. False allows graph breaks (more compatible), True enforces no breaks (maximum optimization but fragile). Default: False")
|
|
parser.add_argument("--compile_dynamic", action="store_true",
|
|
help="Handle varying input shapes without recompilation. False specializes for exact shapes, True creates dynamic kernels. Default: False")
|
|
parser.add_argument("--compile_dynamo_cache_size_limit", type=int, default=64,
|
|
help="Max cached compiled versions per function. Higher = more memory, lower = more recompilation. Default: 64")
|
|
parser.add_argument("--compile_dynamo_recompile_limit", type=int, default=128,
|
|
help="Max recompilation attempts before falling back to uncompiled. Safety limit to prevent infinite loops. Default: 128")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
"""Main CLI function"""
|
|
debug.log(f"SeedVR2 Video Upscaler CLI started at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", category="dit", force=True)
|
|
|
|
# Parse arguments
|
|
args = parse_arguments()
|
|
debug.enabled = args.debug
|
|
|
|
debug.log("Arguments:", category="setup")
|
|
for key, value in vars(args).items():
|
|
debug.log(f" {key}: {value}", category="none")
|
|
|
|
if args.vae_encode_tiling_enabled and (args.vae_encode_tile_overlap[0] >= args.vae_encode_tile_size[0] or args.vae_encode_tile_overlap[1] >= args.vae_encode_tile_size[1]):
|
|
debug.log(f"VAE encode tile overlap {args.vae_encode_tile_overlap} must be smaller than tile size {args.vae_encode_tile_size}", level="ERROR", category="vae", force=True)
|
|
sys.exit(1)
|
|
|
|
if args.vae_decode_tiling_enabled and (args.vae_decode_tile_overlap[0] >= args.vae_decode_tile_size[0] or args.vae_decode_tile_overlap[1] >= args.vae_decode_tile_size[1]):
|
|
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)
|
|
|
|
if args.debug:
|
|
if platform.system() == "Darwin":
|
|
debug.log("You are running on macOS and will use the MPS backend!", category="info", force=True)
|
|
else:
|
|
# Show actual CUDA device visibility
|
|
debug.log(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set (all)')}", category="device")
|
|
if torch.cuda.is_available():
|
|
debug.log(f"torch.cuda.device_count(): {torch.cuda.device_count()}", category="device")
|
|
debug.log(f"Using device index 0 inside script (mapped to selected GPU)", category="device")
|
|
|
|
try:
|
|
# Ensure --output is a directory when using PNG format
|
|
if args.output_format == "png":
|
|
output_path_obj = Path(args.output)
|
|
if output_path_obj.suffix: # an extension is present, strip it
|
|
args.output = str(output_path_obj.with_suffix(''))
|
|
|
|
debug.log(f"Output will be saved to: {args.output}", category="file")
|
|
|
|
# Extract frames from video
|
|
debug.log(f"Extracting frames from video...", category="generation")
|
|
start_time = time.time()
|
|
frames_tensor, original_fps = extract_frames_from_video(
|
|
args.video_path,
|
|
args.skip_first_frames,
|
|
args.load_cap,
|
|
args.prepend_frames
|
|
)
|
|
|
|
debug.log(f"Frame extraction time: {time.time() - start_time:.2f}s", category="timing")
|
|
# debug.log(f"Initial VRAM: {torch.cuda.memory_allocated() / 1024**3:.2f}GB", category="memory")
|
|
|
|
# Parse GPU list
|
|
if platform.system() == "Darwin":
|
|
device_list = ["0"]
|
|
else:
|
|
device_list = [d.strip() for d in str(args.cuda_device).split(',') if d.strip()] if args.cuda_device else ["0"]
|
|
|
|
if args.debug:
|
|
debug.log(f"Using devices: {device_list}", category="device")
|
|
processing_start = time.time()
|
|
# Download both DiT and VAE models
|
|
if not download_weight(dit_model=args.model, vae_model=DEFAULT_VAE, model_dir=args.model_dir, debug=debug):
|
|
debug.log("Failed to download required models. Check console output above.", level="ERROR", category="download", force=True)
|
|
sys.exit(1)
|
|
result = _gpu_processing(frames_tensor, device_list, args)
|
|
generation_time = time.time() - processing_start
|
|
|
|
debug.log(f"Generation time: {generation_time:.2f}s", category="timing")
|
|
if platform.system() != "Darwin":
|
|
debug.log(f"Peak VRAM usage: {torch.cuda.max_memory_allocated() / 1024**3:.2f}GB", category="memory")
|
|
|
|
if args.temporal_overlap > 0:
|
|
debug.log(f"Applying temporal overlap with blending", category="generation")
|
|
result = apply_temporal_overlap_blending(result, args.batch_size, args.temporal_overlap)
|
|
debug.log(f"Result shape: {result.shape}, dtype: {result.dtype}", category="memory")
|
|
|
|
if args.prepend_frames > 0:
|
|
debug.log(f"Removing prepended ({args.prepend_frames}) frames from the results)", category="generation")
|
|
result = result[args.prepend_frames:]
|
|
debug.log(f"Result shape after removing prepended frames: {result.shape}", category="info")
|
|
|
|
# After generation_time calculation, choose saving method
|
|
if args.output_format == "png":
|
|
# Ensure output treated as directory
|
|
output_dir = args.output
|
|
base_name = Path(args.video_path).stem + "_upscaled"
|
|
debug.log(f"Saving PNG frames to directory: {output_dir}", category="file")
|
|
|
|
save_start = time.time()
|
|
save_frames_to_png(result, output_dir, base_name)
|
|
|
|
debug.log(f"Save time: {time.time() - save_start:.2f}s", category="timing")
|
|
|
|
else:
|
|
# Save video
|
|
debug.log(f"Saving upscaled video to: {args.output}", category="file")
|
|
save_start = time.time()
|
|
save_frames_to_video(result, args.output, original_fps)
|
|
debug.log(f"Save time: {time.time() - save_start:.2f}s", category="timing")
|
|
|
|
total_time = time.time() - start_time
|
|
debug.log(f"Upscaling completed successfully!", category="success", force=True)
|
|
if args.output_format == "png":
|
|
debug.log(f"PNG frames saved in directory: {args.output}", category="file", force=True)
|
|
else:
|
|
debug.log(f"Output saved to video: {args.output}", category="file", force=True)
|
|
debug.log(f"Total processing time: {total_time:.2f}s", category="timing", force=True)
|
|
debug.log(f"Average FPS: {len(frames_tensor) / generation_time:.2f} frames/sec", category="timing", force=True)
|
|
|
|
except Exception as e:
|
|
debug.log(f"Error during processing: {e}", level="ERROR", category="generation", force=True)
|
|
import traceback
|
|
traceback.print_exc()
|
|
sys.exit(1)
|
|
|
|
finally:
|
|
debug.log(f"Process {os.getpid()} terminating - VRAM will be automatically freed", category="cleanup", force=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |