bnair/VMAFOptimiser
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added actualy HW encode

Browse Source
This commit is contained in:
bnair123
2025-12-31 23:39:21 +04:00
parent 05a306dc42
commit b4a82e0db5
3 changed files with 218 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

139
optimize_library.py
View File

@@ -10,7 +10,7 @@ import signal
from pathlib import Path
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed, ProcessPoolExecutor
from threading import Lock
from threading import Lock, get_ident
# --- Configuration ---
DEFAULT_VMAF = 95.0
@@ -23,11 +23,27 @@ TARGETS = [94.0, 93.0, 92.0, 90.0]
MIN_SAVINGS_PERCENT = 12.0
TARGET_SAVINGS_FOR_ESTIMATE = 15.0
# Hardware encoder quality compensation (GPU needs higher VMAF target to match CPU quality)
HW_ENCODER_VMAF_OFFSET = 2.0
HW_ENCODERS = {"av1_amf", "av1_nvenc", "av1_qsv", "av1_vaapi"}
# Global state for resume capability
_processed_files = set()
_lock = Lock()
_shutdown_requested = False
_AB_AV1_HELP_CACHE = {}
_hw_worker_id = None # Thread ID of the designated hardware worker
def claim_hardware_worker():
"""Attempt to claim hardware worker status for this thread. Returns True if claimed."""
global _hw_worker_id
with _lock:
thread_id = get_ident()
if _hw_worker_id is None:
_hw_worker_id = thread_id
return True
return _hw_worker_id == thread_id
def signal_handler(signum, frame):
@@ -202,15 +218,30 @@ def run_command_streaming(cmd, description=""):
return process.returncode
def run_crf_search(filepath, target_vmaf, preset, temp_dir, use_hw=False, hwaccel=None):
def run_crf_search(
filepath,
target_vmaf,
preset,
temp_dir,
encoder="svt-av1",
use_hw=False,
hwaccel=None,
):
"""Run CRF search for a specific VMAF target"""
is_hw_encoder = encoder in HW_ENCODERS
# Apply VMAF offset for hardware encoders to match software quality
effective_vmaf = (
target_vmaf + HW_ENCODER_VMAF_OFFSET if is_hw_encoder else target_vmaf
)
cmd = [
"ab-av1",
"crf-search",
"-i",
str(filepath),
"--min-vmaf",
str(target_vmaf),
str(effective_vmaf),
"--preset",
str(preset),
"--max-encoded-percent",
@@ -218,18 +249,24 @@ def run_crf_search(filepath, target_vmaf, preset, temp_dir, use_hw=False, hwacce
"--temp-dir",
temp_dir,
"--samples",
"4", # Use 4 samples for speed/accuracy balance
"4",
]
# Hardware encoding support
# Add encoder if not default
if encoder != "svt-av1":
if ab_av1_supports("crf-search", "--encoder"):
cmd.extend(["--encoder", encoder])
# Hardware decode acceleration
if use_hw and hwaccel:
if ab_av1_supports("crf-search", "--enc-input"):
cmd.extend(["--enc-input", f"hwaccel={hwaccel}"])
if hwaccel == "vaapi":
cmd.extend(["--enc-input", "hwaccel_output_format=vaapi"])
print(f" - Searching for CRF to hit VMAF {target_vmaf}...")
returncode = run_command_streaming(cmd, f"crf-search VMAF {target_vmaf}")
vmaf_label = f"VMAF {effective_vmaf}" if is_hw_encoder else f"VMAF {target_vmaf}"
print(f" - Searching for CRF to hit {vmaf_label}...")
returncode = run_command_streaming(cmd, f"crf-search {vmaf_label}")
if returncode == 0:
# Parse output to find CRF and predicted size
@@ -274,7 +311,13 @@ def run_crf_search(filepath, target_vmaf, preset, temp_dir, use_hw=False, hwacce
def find_target_savings_params(
filepath, start_vmaf, preset, temp_dir, use_hw=False, hwaccel=None
filepath,
start_vmaf,
preset,
temp_dir,
encoder="svt-av1",
use_hw=False,
hwaccel=None,
):
"""Find VMAF target that achieves minimum savings"""
print(f"\n --- Finding VMAF for {TARGET_SAVINGS_FOR_ESTIMATE}% savings ---")
@@ -288,7 +331,9 @@ def find_target_savings_params(
print(
f" Testing VMAF {target} for {TARGET_SAVINGS_FOR_ESTIMATE}% target... (test {i + 1}/{len(test_targets)})"
)
result = run_crf_search(filepath, target, preset, temp_dir, use_hw, hwaccel)
result = run_crf_search(
filepath, target, preset, temp_dir, encoder, use_hw, hwaccel
)
if result:
predicted_savings = 100.0 - result["predicted_percent"]
@@ -317,7 +362,9 @@ def find_target_savings_params(
return None
def run_encode(filepath, output_path, crf, preset, use_hw=False, hwaccel=None):
def run_encode(
filepath, output_path, crf, preset, encoder="svt-av1", use_hw=False, hwaccel=None
):
"""Run full encoding with real-time output streaming"""
cmd = [
"ab-av1",
@@ -334,7 +381,10 @@ def run_encode(filepath, output_path, crf, preset, use_hw=False, hwaccel=None):
"copy",
]
# Hardware encoding support
if encoder != "svt-av1":
if ab_av1_supports("encode", "--encoder"):
cmd.extend(["--encoder", encoder])
if use_hw and hwaccel:
if ab_av1_supports("encode", "--enc-input"):
cmd.extend(["--enc-input", f"hwaccel={hwaccel}"])
@@ -401,10 +451,26 @@ def refresh_plex(plex_url, plex_token):
print(f" ⚠️ Failed to refresh Plex: {e}")
def process_file(filepath, log_dir, log_name, preset, use_hw=False, hwaccel=None):
def process_file(
filepath,
log_dir,
log_name,
preset,
hw_encoder="svt-av1",
use_hw_mode=False,
hwaccel=None,
):
"""Process a single video file with intelligent VMAF targeting"""
global _shutdown_requested
# Determine if THIS worker should use hardware encoder
use_hw = False
if use_hw_mode and hwaccel and hw_encoder in HW_ENCODERS:
use_hw = claim_hardware_worker()
# HW worker uses hardware encoder; CPU workers use svt-av1
encoder = hw_encoder if use_hw else "svt-av1"
filepath = Path(filepath)
lock_file = Path(log_dir).parent / ".lock" / f"{filepath.name}.lock"
@@ -449,21 +515,21 @@ def process_file(filepath, log_dir, log_name, preset, use_hw=False, hwaccel=None
# Step 1: Try VMAF 94
print(f"\n [Step 1] Testing VMAF 94...")
search_result_94 = run_crf_search(
filepath, 94.0, preset, str(temp_dir), use_hw, hwaccel
filepath, 94.0, preset, str(temp_dir), encoder, use_hw, hwaccel
)
if not search_result_94:
print(f" !! Could not hit VMAF 94")
search_result_94 = run_crf_search(
filepath, 93.0, preset, str(temp_dir), use_hw, hwaccel
filepath, 93.0, preset, str(temp_dir), encoder, use_hw, hwaccel
)
if not search_result_94:
search_result_94 = run_crf_search(
filepath, 92.0, preset, str(temp_dir), use_hw, hwaccel
filepath, 92.0, preset, str(temp_dir), encoder, use_hw, hwaccel
)
if not search_result_94:
search_result_94 = run_crf_search(
filepath, 90.0, preset, str(temp_dir), use_hw, hwaccel
filepath, 90.0, preset, str(temp_dir), encoder, use_hw, hwaccel
)
if not search_result_94:
@@ -499,7 +565,7 @@ def process_file(filepath, log_dir, log_name, preset, use_hw=False, hwaccel=None
)
search_result_93 = run_crf_search(
filepath, 93.0, preset, str(temp_dir), use_hw, hwaccel
filepath, 93.0, preset, str(temp_dir), encoder, use_hw, hwaccel
)
if search_result_93:
predicted_savings_93 = 100.0 - search_result_93["predicted_percent"]
@@ -524,7 +590,7 @@ def process_file(filepath, log_dir, log_name, preset, use_hw=False, hwaccel=None
f" → Finding VMAF for {TARGET_SAVINGS_FOR_ESTIMATE}% savings..."
)
target_result = find_target_savings_params(
filepath, 93.0, preset, str(temp_dir), use_hw, hwaccel
filepath, 93.0, preset, str(temp_dir), encoder, use_hw, hwaccel
)
provide_recommendations(
@@ -559,7 +625,13 @@ def process_file(filepath, log_dir, log_name, preset, use_hw=False, hwaccel=None
start_time = time.time()
res = run_encode(
filepath, temp_output, encode_params["crf"], preset, use_hw, hwaccel
filepath,
temp_output,
encode_params["crf"],
preset,
encoder,
use_hw,
hwaccel,
)
if res != 0:
@@ -688,10 +760,18 @@ def main():
"Examples: auto, vaapi, d3d11va, videotoolbox. Use 'none' to disable."
),
)
parser.add_argument(
"--encoder",
default="svt-av1",
help=(
"Video encoder to use. Default: svt-av1 (CPU). "
"Hardware encoders: av1_amf (AMD), av1_nvenc (NVIDIA), av1_qsv (Intel)."
),
)
parser.add_argument(
"--use-hardware-worker",
action="store_true",
help="Use 1 hardware encoding worker + rest CPU workers (requires --hwaccel)",
help="Use 1 hardware encoding worker + rest CPU workers (requires --encoder with HW encoder)",
)
parser.add_argument(
"--plex-url",
@@ -763,7 +843,9 @@ def main():
fail_count = 0
# Hardware worker configuration
use_hw_primary = args.use_hardware_worker and hwaccel is not None
# HW worker uses specified encoder; CPU workers use svt-av1
hw_encoder = args.encoder if args.encoder in HW_ENCODERS else None
use_hw_primary = args.use_hardware_worker and hw_encoder is not None
if args.workers == 1:
# Single thread - just process files
@@ -772,7 +854,13 @@ def main():
break
processed_count += 1
result = process_file(
file_path, args.log_dir, log_name, args.preset, use_hw_primary, hwaccel
file_path,
args.log_dir,
log_name,
args.preset,
args.encoder,
use_hw_primary,
hwaccel,
)
if result:
success_count += 1
@@ -786,15 +874,16 @@ def main():
if _shutdown_requested:
break
# Use hardware for first file, CPU for rest
use_hw_for_this = use_hw_primary and len(futures) == 0
# All workers try to claim HW; only the first thread succeeds
# and will use HW for ALL its tasks
future = executor.submit(
process_file,
file_path,
args.log_dir,
log_name,
args.preset,
use_hw_for_this,
args.encoder,
use_hw_primary,
hwaccel,
)
futures.append(future)