-
Notifications
You must be signed in to change notification settings - Fork 1
/
Av1anStaxRipWrapperRav1e.py
188 lines (160 loc) · 9.27 KB
/
Av1anStaxRipWrapperRav1e.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
import argparse
import subprocess
import sys
sys.stdout.reconfigure(encoding='utf-8')
# This script is more specialized for rav1e.
# Use the generic script for other encoders. https://github.com/Kidsnd274/Av1anStaxRipWrapper
# Functions
def add_argument(curr, new):
return_string = curr
if curr == "":
return_string = new
else:
return_string += (" " + new)
return return_string
def set_path(path):
import os
import pathlib
staxrip_path = pathlib.Path(path)
av1an_path = staxrip_path / "Apps" / "Encoders" / "Av1an"
rav1e_path = staxrip_path / "Apps" / "Encoders" / "rav1e"
vp_path = staxrip_path / "Apps" / "Encoders" / "Av1anStaxRipWrapper" / "VapourSynth"
environ = os.environ
environ["PATH"] = f"{str(av1an_path)};{str(rav1e_path)};{str(vp_path)};{environ['PATH']}"
return environ
def print_welcome():
print("=================================================")
print("Av1anStaxRipWrapperRav1e")
print("https://github.com/Kidsnd274/Av1anStaxRipWrapper\n")
print("This script is more specialized for rav1e")
print("Use the generic script for other encoders")
print("=================================================")
print("")
def print_version(parser_args):
if parser_args.staxrip_startup_dir is not None:
my_env = set_path(parser_args.staxrip_startup_dir)
else:
import os
my_env = os.environ
try:
subprocess.run("ffmpeg -version", shell=False, env=my_env)
except FileNotFoundError:
print("ffmpeg not found!")
print("\n--------------------------------\n")
try:
subprocess.run("av1an --version", shell=False, env=my_env)
except FileNotFoundError:
print("Av1an not found!")
print("\n--------------------------------\n")
try:
subprocess.run("rav1e --version", shell=False, env=my_env)
except FileNotFoundError:
print("rav1e not found!")
print("\n--------------------------------\n")
exit(0)
# Command Line Arguments
parser = argparse.ArgumentParser(description="Av1an wrapper for StaxRip")
parser.add_argument('--version', action='store_true', help="Print Av1an, ffmpeg and rav1e versions")
parser.add_argument('-i', dest="input", type=str, help="Input File (for StaxRip)")
parser.add_argument('-o', dest="output", type=str, help="Output File (for StaxRip)")
parser.add_argument('-t', dest="tempdir", type=str, help="Temp Directory (for StaxRip)")
parser.add_argument('-s', '--staxrip-startup-dir', dest="staxrip_startup_dir", type=str, required=False, help="Specify StaxRip Startup Directory so that the wrapper script will automatically add important folders to PATH for av1an to detect (only needed for portable installations)")
parser.add_argument('--photon-noise', dest="photon_noise", type=str, required=False, help="Generates a photon noise table and applies it using grain synthesis [strength: 0-64] (disabled by default) (Av1an parameter)")
parser.add_argument('--chroma-noise', dest="chroma_noise", action='store_true', help="Adds chroma grain synthesis to the grain table generated by `--photon-noise`. (Default: false) (Av1an parameter)")
parser.add_argument('--sc-downscale-height', dest="sc_downscale_height", type=str, required=False, help="Optional downscaling for scene detection. By default, no downscaling is performed. (Av1an parameter)")
parser.add_argument('--pix-format', dest="pix_format", type=str, required=False, help="FFmpeg pixel format")
# Threading Arguments (do not specify these commands if you want to use Automatic Thread Detection)
parser.add_argument('--workers', type=str, required=False, help="Number of workers to spawn [0 = automatic] (Av1an Paramter)")
parser.add_argument('--set-thread-affinity', dest="set_thread_affinity", type=str, required=False, help="Pin each worker to a specific set of threads of this size (disabled by default) (Av1an parameter)")
parser.add_argument('--disable-automatic-thread-detection', dest="disable_automatic_thread_detection", action='store_true', help="Disable the wrapper's automatic thread detection")
# Rav1e arguments
parser.add_argument('--quantizer', type=str, required=False, help="Quantizer (0-255), smaller values are higher quality (default: 100) (rav1e parameter)") # Quantizer (0-255), smaller values are higher quality (default: 100)
parser.add_argument('--speed', type=str, required=False, help="Speed level (0 is best quality, 10 is fastest)\nSpeeds 10 and 0 are extremes and are generally not recommended\n[default: 6] (rav1e parameter)") # Speed level 0-10 (0 is best quality, 10 is fastest) (default: 6)
parser.add_argument('--tiles', type=str, required=False, help="Number of tiles. Tile-cols and tile-rows are overridden so that the video has at least this many tiles (rav1e parameter)")
parser.add_argument('--threads', type=str, required=False, help="Set the threadpool size. If 0, will use the number of logical CPUs. rav1e will use up to this many threads.\nAdditional tiles may be needed to increase thread utilization\n[default: 0] (rav1e parameter)")
parser_args = parser.parse_args()
print_welcome()
if parser_args.version:
print_version(parser_args)
if parser_args.input is None or parser_args.output is None or parser_args.tempdir is None:
print("The arguments, -i, -o, -t are required to work!")
print("Run --help for more information")
exit(1)
input_file = parser_args.input
output_file = parser_args.output
tempdir = parser_args.tempdir
# # Parsing rav1e arguments
rav1e_argument_string = ""
if parser_args.speed is not None:
rav1e_argument_string = add_argument(rav1e_argument_string, f"--speed {parser_args.speed}")
if parser_args.quantizer is not None:
rav1e_argument_string = add_argument(rav1e_argument_string, f"--quantizer {parser_args.quantizer}")
if parser_args.tiles is not None:
rav1e_argument_string = add_argument(rav1e_argument_string, f"--tiles {parser_args.tiles}")
if parser_args.threads is not None:
rav1e_argument_string = add_argument(rav1e_argument_string, f"--threads {parser_args.threads}")
# Automatic Thread Detection
thread_detection = False
if not parser_args.disable_automatic_thread_detection and parser_args.workers is None and parser_args.set_thread_affinity is None:
thread_detection = True
if thread_detection: # Checking for new Intel architecture
import psutil
logical_count = psutil.cpu_count(logical = True)
physical_count = psutil.cpu_count(logical = False)
if (logical_count / physical_count) % 1 != 0:
thread_detection = False # Intel CPU detected
print("New Intel CPU architecture with performance and efficiency cores detected!\nNot passing thread detection to av1an...\n")
if thread_detection: # Checking for Hyperthreading or SMT
import psutil
logical_count = psutil.cpu_count(logical = True)
physical_count = psutil.cpu_count(logical = False)
if (logical_count / physical_count) == 2:
hyperthreading = True
else:
hyperthreading = False
if thread_detection: # Setting values
if hyperthreading:
cpu_workers = physical_count
cpu_thread_affinity = 2
else:
cpu_workers = physical_count
cpu_thread_affinity = 1
print(f"THREADING INFORMATION:\n Hyperthreading / SMT: {hyperthreading}\n Workers: {cpu_workers}\n Thread Affinity: {cpu_thread_affinity}\n\n")
else:
print("THREADING INFORMATION:\n Automatic Thread Detection: DISABLED\n\n")
# If StaxRip path given, automatically add important folders to PATH
if parser_args.staxrip_startup_dir is not None:
my_env = set_path(parser_args.staxrip_startup_dir)
av1an_exec = "av1an.exe"
command = av1an_exec
command = add_argument(command, "--verbose -y --resume -a=\"-an\" -e rav1e")
# Thread arguments
if thread_detection:
command = add_argument(command, f"--workers {cpu_workers} --set-thread-affinity {cpu_thread_affinity}")
else:
if parser_args.workers is not None:
command = add_argument(command, f"--workers {parser_args.workers}")
if parser_args.set_thread_affinity is not None:
command = add_argument(command, f"--set-thread-affinity {parser_args.set_thread_affinity}")
if parser_args.photon_noise is not None:
command = add_argument(command, f"--photon-noise {parser_args.photon_noise}")
if parser_args.chroma_noise:
command = add_argument(command, f"--chroma-noise")
if parser_args.sc_downscale_height is not None:
command = add_argument(command, f"--sc-downscale-height {parser_args.sc_downscale_height}")
if parser_args.pix_format is not None:
command = add_argument(command, f"--pix-format {parser_args.pix_format}")
if rav1e_argument_string != "":
command = add_argument(command, f"-v=\"{rav1e_argument_string} --no-scene-detection\"")
command = add_argument(command, f"-i \"{input_file}\" -o \"{output_file}\" --temp \"{tempdir}\"")
sys.stdout.write("Starting av1an... Check new console window for progress\n")
sys.stdout.write("Command: " + str(command) + "\n")
sys.stdout.flush()
if parser_args.staxrip_startup_dir is not None:
process = subprocess.run(command, shell=False, creationflags=subprocess.CREATE_NEW_CONSOLE, env=my_env)
else:
process = subprocess.run(command, shell=False, creationflags=subprocess.CREATE_NEW_CONSOLE) # Assume everything is in PATH
if process.returncode != 0:
print(process.stderr)
print("Error occurred when transcoding with av1an. Check logs")
exit(1)