yt_common: automation: audio handling updates

.gitignore

@@ -22,3 +22,4 @@ ffmpeg2pass*.log
 **/results/
 **/__pycache__/
 **/*.egg-info/
+**/bdmv/

Vivy/01/01.vpy

@@ -50,7 +50,7 @@ LETTERBOX: List[Range] = [
 
 
 def filter_basic() -> vs.VideoNode:
-    wakas, ref = SOURCE.source()
+    wakas, ref = SOURCE.dhs_source()
     wakas = [waka[:37301] + core.std.BlankClip(waka, length=6) + waka[37301:] for waka in wakas]
     waka = wakas[0]
     waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE)
@@ -71,6 +71,6 @@ def filter() -> vs.VideoNode:
 
 
 if __name__ == "__main__":
-    SelfRunner(CONFIG, filter, filter_basic)
+    SelfRunner(CONFIG, SOURCE, filter, filter_basic)
 else:
     filter()

Vivy/02/02.vpy

@@ -58,7 +58,7 @@ NOAA: List[Range] = PIXELSHIT
 
 
 def filter_basic() -> vs.VideoNode:
-    wakas, ref = SOURCE.source()
+    wakas, ref = SOURCE.dhs_source()
     wakas = [waka[:32344] + waka[32349:] for waka in wakas]
     waka = wakas[0]
     waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE)
@@ -78,6 +78,6 @@ def filter() -> vs.VideoNode:
 
 
 if __name__ == "__main__":
-    SelfRunner(CONFIG, filter, filter_basic)
+    SelfRunner(CONFIG, SOURCE, filter, filter_basic)
 else:
     filter()

Vivy/03/03.vpy

@@ -42,7 +42,7 @@ LETTERBOX: List[Range] = [(0, 432)]
 
 
 def filter_basic() -> vs.VideoNode:
-    wakas, ref = SOURCE.source()
+    wakas, ref = SOURCE.dhs_source()
     waka = wakas[0]
     waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE)
     src = bounded_dehardsub(waka, ref, SIGNS_RU)
@@ -62,6 +62,6 @@ def filter() -> vs.VideoNode:
 
 
 if __name__ == "__main__":
-    SelfRunner(CONFIG, filter, filter_basic)
+    SelfRunner(CONFIG, SOURCE, filter, filter_basic)
 else:
     filter()

Vivy/04/04.vpy

@@ -14,8 +14,6 @@ from typing import List
 import os
 
 core = vs.core
-core.num_threads = 16
-
 
 EPNUM: int = int(os.path.basename(os.path.splitext(__file__)[0]))
 CONFIG: VivyConfig = VivyConfig(EPNUM)
@@ -39,7 +37,7 @@ LETTERBOX: List[Range] = [(0, 1151)]
 
 
 def filter_basic() -> vs.VideoNode:
-    wakas, ref = SOURCE.source()
+    wakas, ref = SOURCE.dhs_source()
     wakas = [w[:33665] + core.std.BlankClip(w, length=21) + w[33665:] for w in wakas]
     waka = wakas[0]
     waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE)
@@ -62,6 +60,6 @@ def filter() -> vs.VideoNode:
 
 
 if __name__ == "__main__":
-    SelfRunner(CONFIG, filter, filter_basic)
+    SelfRunner(CONFIG, SOURCE, filter, filter_basic)
 else:
     filter()

Vivy/05/05.vpy

@@ -50,7 +50,7 @@ LETTERBOX: List[Range] = [(0, 1153)]
 
 
 def filter_basic() -> vs.VideoNode:
-    wakas, ref = SOURCE.source()
+    wakas, ref = SOURCE.dhs_source()
     waka = wakas[0]
     waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE)
     src = bounded_dehardsub(waka, ref, SIGNS_RU, wakas)
@@ -71,6 +71,6 @@ def filter() -> vs.VideoNode:
 
 
 if __name__ == "__main__":
-    SelfRunner(CONFIG, filter, filter_basic)
+    SelfRunner(CONFIG, SOURCE, filter, filter_basic)
 else:
     filter()

Vivy/06/06.vpy

@@ -110,7 +110,7 @@ LETTERBOX: List[Range] = [(0, 2150), (8791, 10693), (13427, 15153), (27878, 2800
 
 
 def filter_basic() -> vs.VideoNode:
-    wakas, ref = SOURCE.source()
+    wakas, ref = SOURCE.dhs_source()
     wakas = [w[:33669] + w.std.BlankClip(length=17) + w[33669:] for w in wakas]
     ref = ref.resize.Point(range_in=CRange.FULL, range=CRange.LIMITED) if SOURCE.ref_is_funi else ref
     waka = wakas[0]
@@ -134,6 +134,6 @@ def filter() -> vs.VideoNode:
 
 
 if __name__ == "__main__":
-    SelfRunner(CONFIG, filter, filter_basic)
+    SelfRunner(CONFIG, SOURCE, filter, filter_basic)
 else:
     filter()

Vivy/07/07.vpy

@@ -91,7 +91,7 @@ LETTERBOX: List[Range] = []
 
 
 def filter_basic() -> vs.VideoNode:
-    wakas, ref = SOURCE.source()
+    wakas, ref = SOURCE.dhs_source()
     waka = wakas[0]
     waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE)
     src = bounded_dehardsub(waka, ref, SIGNS_RU, wakas)
@@ -116,6 +116,6 @@ def filter() -> vs.VideoNode:
 
 
 if __name__ == "__main__":
-    SelfRunner(CONFIG, filter, filter_basic)
+    SelfRunner(CONFIG, SOURCE, filter, filter_basic)
 else:
     filter()

Vivy/setup.py

@@ -19,7 +19,7 @@ setuptools.setup(
         "Operating System :: OS Independent",
     ],
     package_data={
-        'vivy_common': ['py.typed', 'workraw-settings', 'final-settings', 'shaders/FSRCNNX_x2_56-16-4-1.glsl'],
+        'vivy_common': ['py.typed', 'workraw-settings', 'final-settings'],
     },
     python_requires='>=3.8',
 )

Vivy/vivy_common/filter.py

@@ -12,16 +12,14 @@ from lvsfunc.types import Range
 from typing import List, Optional, Union
 
 from yt_common import antialiasing
+from yt_common.data import FSRCNNX
 from yt_common.denoise import bm3d
 from yt_common.deband import morpho_mask
 
-import os
 import vsutil
 
 core = vs.core
 
-FSRCNNX = os.path.join(os.path.dirname(__file__), "shaders/FSRCNNX_x2_56-16-4-1.glsl")
-
 
 def fsrcnnx_rescale(src: vs.VideoNode, noscale: Optional[List[Range]] = None,
                     kernel: Optional[lvf.kernels.Kernel] = None) -> vs.VideoNode:

yt_common/setup.py

@@ -19,7 +19,7 @@ setuptools.setup(
         "Operating System :: OS Independent",
     ],
     package_data={
-        'yt_common': ['py.typed'],
+        'yt_common': ['py.typed', 'shaders/FSRCNNX_x2_56-16-4-1.glsl'],
     },
     python_requires='>=3.8',
 )

yt_common/yt_common/__init__.py

@@ -1 +1 @@
-from . import antialiasing, automation, config, deband, denoise, logging, source  # noqa: F401
+from . import antialiasing, automation, config, data, deband, denoise, logging, source  # noqa: F401

yt_common/yt_common/antialiasing.py

@@ -35,9 +35,9 @@ def combine_mask(clip: vs.VideoNode, weak: Union[Range, List[Range], None] = Non
 
 
 def sraa_clamp(clip: vs.VideoNode, mask: Optional[vs.VideoNode] = None,
-               strength: float = 2, opencl: bool = False,
+               strength: float = 2, opencl: bool = True,
                postprocess: Optional[Callable[[vs.VideoNode], vs.VideoNode]] = None) -> vs.VideoNode:
-    sraa = upscaled_sraa(clip, rfactor=1.3, opencl=opencl, downscaler=Bicubic(b=0, c=1/2).scale)
+    sraa = upscaled_sraa(clip, rfactor=1.3, nnedi3cl=opencl, downscaler=Bicubic(b=0, c=1/2).scale)
     sraa = postprocess(sraa) if postprocess else sraa
     clamp = clamp_aa(clip, nnedi3(clip, opencl=opencl), sraa, strength=strength)
     return core.std.MaskedMerge(clip, clamp, mask, planes=0) if mask else clamp

yt_common/yt_common/automation.py

@@ -8,16 +8,15 @@ import shutil
 import string
 import subprocess
 
-from typing import Any, BinaryIO, Callable, List, Optional, Sequence, Union, cast
+from typing import Any, BinaryIO, Callable, List, Optional, Sequence, cast
 
 from .config import Config
 from .logging import log
-from .source import AMAZON_FILENAME_CBR, AMAZON_FILENAME_VBR, ER_FILENAME, SUBSPLS_FILENAME, FUNI_INTRO, glob_filename
+from .source import FileSource
 
 core = vs.core
 
-AUDIO_OVERRIDE: str = "audio.mka"
+AUDIO_ENCODE: str = "_audiogetter_encode.mka"
-AUDIO_CUT: str = "_audiogetter_cut.mka"
 
 
 def bin_to_plat(binary: str) -> str:
@@ -109,6 +108,7 @@ class AudioGetter():
     TODO: really should modularize this a bit instead of assuming amazon->funi
     """
     config: Config
+    src: FileSource
 
     audio_file: str
     audio_start: int
@@ -116,77 +116,41 @@ class AudioGetter():
 
     cleanup: List[str]
 
-    def __init__(self, config: Config, override: Optional[str] = None) -> None:
+    def __init__(self, config: Config, src: FileSource) -> None:
         self.config = config
+        self.src = src
 
         self.audio_start = 0
         self.video_src = None
         self.cleanup = []
 
-        if override is not None:
+    def trim_audio(self, ftrim: Optional[acsuite.types.Trim] = None) -> str:
-            if os.path.isfile(override):
+        trims = self.src.get_audio()
-                self.audio_file = override
+        if not trims or len(trims) > 1:
-            else:
+            raise NotImplementedError("Please implement multifile trimming")
-                raise FileNotFoundError(f"Audio file {override} not found!")
+        audio_cut = acsuite.eztrim(trims[0].path, cast(acsuite.types.Trim, trims[0].trim))[0]
+        self.cleanup.append(audio_cut)
 
-        # drop "audio.m4a" into the folder and it'll get used
+        if ftrim:
-        if os.path.isfile(AUDIO_OVERRIDE):
+            audio_cut = acsuite.eztrim(audio_cut, ftrim, ref_clip=self.src.source())[0]
-            self.audio_file = AUDIO_OVERRIDE
+            self.cleanup.append(audio_cut)
-            return
 
-        # look for amazon first
+        return audio_cut
-        if os.path.isfile(self.config.format_filename(AMAZON_FILENAME_CBR)):
-            self.audio_file = self.config.format_filename(AMAZON_FILENAME_CBR)
-            self.video_src = core.ffms2.Source(self.audio_file)
-            log.success("Found Amazon audio")
-            return
 
-        if os.path.isfile(self.config.format_filename(AMAZON_FILENAME_VBR)):
+    def encode_audio(self, path: str, args: List[str]) -> str:
-            self.audio_file = self.config.format_filename(AMAZON_FILENAME_VBR)
+        ffmpeg_args = [
-            self.video_src = core.ffms2.Source(self.audio_file)
+            "ffmpeg",
-            log.success("Found Amazon audio")
+            "-hide_banner", "-loglevel", "panic",
-            return
+            "-i", path,
+            "-y",
+            "-map", "0:a",
+        ] + args + [AUDIO_ENCODE]
+        print("+ " + " ".join(ffmpeg_args))
+        subprocess.call(ffmpeg_args)
 
-        try:
+        self.cleanup.append(AUDIO_ENCODE)
-            self.audio_file = glob_filename(self.config.format_filename(SUBSPLS_FILENAME))
-            self.video_src = core.ffms2.Source(self.audio_file)
-        except FileNotFoundError:
-            pass
-        try:
-            self.audio_file = glob_filename(self.config.format_filename(ER_FILENAME))
-            self.video_src = core.ffms2.Source(self.audio_file)
-        except FileNotFoundError:
-            log.error("Could not find audio")
-            raise
 
-        self.audio_start = FUNI_INTRO
+        return AUDIO_ENCODE
-        log.warn("No Amazon audio, falling back to Funi")
-
-    def trim_audio(self, src: vs.VideoNode,
-                   trims: Union[acsuite.Trim, List[acsuite.Trim], None] = None) -> str:
-        if isinstance(trims, tuple):
-            trims = [trims]
-
-        if trims is None or len(trims) == 0:
-            if self.audio_start == 0:
-                trims = [(None, None)]
-            else:
-                trims = [(self.audio_start, None)]
-        else:
-            if self.audio_start != 0:
-                trims = [(s+self.audio_start if s is not None and s >= 0 else s,
-                          e+self.audio_start if e is not None and e > 0 else e)
-                         for s, e in trims]
-
-        if os.path.isfile(AUDIO_CUT):
-            os.remove(AUDIO_CUT)
-
-        acsuite.eztrim(self.video_src if self.video_src else src, trims,
-                       self.audio_file, AUDIO_CUT, quiet=True)
-
-        self.cleanup.append(AUDIO_CUT)
-
-        return AUDIO_CUT
 
     def do_cleanup(self) -> None:
         for f in self.cleanup:
@@ -196,6 +160,7 @@ class AudioGetter():
 
 class SelfRunner():
     config: Config
+    src: FileSource
     clip: vs.VideoNode
 
     workraw: bool
@@ -208,9 +173,11 @@ class SelfRunner():
 
     profile: str
 
-    def __init__(self, config: Config, final_filter: Callable[[], vs.VideoNode],
+    def __init__(self, config: Config, source: FileSource, final_filter: Callable[[], vs.VideoNode],
-                 workraw_filter: Optional[Callable[[], vs.VideoNode]] = None) -> None:
+                 workraw_filter: Optional[Callable[[], vs.VideoNode]] = None,
+                 audio_codec: Optional[List[str]] = None) -> None:
         self.config = config
+        self.src = source
         self.video_clean = False
         self.audio_clean = False
 
@@ -222,7 +189,6 @@ class SelfRunner():
         parser.add_argument("-k", "--keep", help="Keep raw video", action="store_true")
         parser.add_argument("-b", "--encoder", type=str, help="Override detected encoder binary.")
         parser.add_argument("-f", "--force", help="Overwrite existing intermediaries.", action="store_true")
-        parser.add_argument("-a", "--audio", type=str, help="Force audio file")
         parser.add_argument("-x", "--suffix", type=str, help="Change the suffix of the mux. \
                             Will be overridden by PROFILE if set.")
         parser.add_argument("-p", "--profile", type=str, help="Set the encoder profile. \
@@ -283,10 +249,13 @@ class SelfRunner():
                                               start, end)
 
         log.status("--- LOOKING FOR AUDIO ---")
-        self.audio = AudioGetter(self.config, args.audio)
+        self.audio = AudioGetter(self.config, self.src)
 
         log.status("--- TRIMMING AUDIO ---")
-        self.audio_file = self.audio.trim_audio(self.clip, (start, end))
+        self.audio_file = self.audio.trim_audio((start, end))
+        if audio_codec:
+            log.status("--- TRANSCODING AUDIO ---")
+            self.audio_file = self.audio.encode_audio(self.audio_file, audio_codec)
 
         try:
             log.status("--- MUXING FILE ---")

yt_common/yt_common/data.py

@@ -0,0 +1,3 @@
+import os
+
+FSRCNNX: str = os.path.join(os.path.dirname(__file__), "shaders/FSRCNNX_x2_56-16-4-1.glsl")

yt_common/yt_common/source.py

@@ -1,15 +1,16 @@
 import vapoursynth as vs
 
-import vsutil
+from acsuite.types import Trim
 
+from lvsfunc.misc import replace_ranges
 from lvsfunc.types import Range
-import lvsfunc as lvf
+from vsutil import depth
 
 import glob
 import os
 
 from abc import ABC, abstractmethod
-from typing import Any, List, Tuple
+from typing import Any, List, NamedTuple, Optional, Tuple, Union
 
 from .config import Config
 from .logging import log
@@ -23,6 +24,23 @@ AMAZON_FILENAME_CBR: str = "{title_long} - {epnum:02d} (Amazon Prime CBR {resolu
 AMAZON_FILENAME_VBR: str = "{title_long} - {epnum:02d} (Amazon Prime VBR {resolution}p).mkv"
 
 
+class FileTrim(NamedTuple):
+    path: str
+    trim: Optional[Trim]
+
+    def apply_trim(self, clip: vs.VideoNode) -> vs.VideoNode:
+        if self.trim is None:
+            return clip
+        s, e = self.trim
+        if s is None and e is None:
+            return clip
+        if s is None:
+            return clip[:e]
+        if e is None:
+            return clip[s:]
+        return clip[s:e]
+
+
 def waka_replace(src: vs.VideoNode, wakas: List[vs.VideoNode], ranges: List[List[Range]]
                  ) -> Tuple[vs.VideoNode, List[vs.VideoNode]]:
     if len(wakas) == 0:
@@ -30,8 +48,8 @@ def waka_replace(src: vs.VideoNode, wakas: List[vs.VideoNode], ranges: List[List
     new_wakas = []
     for waka, r in zip(wakas, ranges):
         tmp = src
-        src = lvf.misc.replace_ranges(src, waka, r)
+        src = replace_ranges(src, waka, r)
-        new_wakas.append(lvf.misc.replace_ranges(waka, tmp, r))
+        new_wakas.append(replace_ranges(waka, tmp, r))
 
     return src, new_wakas
 
@@ -43,7 +61,34 @@ def glob_filename(pattern: str) -> str:
     return res[0]
 
 
-class DehardsubFileFinder(ABC):
+class FileSource(ABC):
+    def _open(self, path: str) -> vs.VideoNode:
+        return depth(core.lsmas.LWLibavSource(path), 16) if path.lower().endswith(".m2ts") \
+            else depth(core.ffms2.Source(path), 16)
+
+    @abstractmethod
+    def get_audio(self) -> List[FileTrim]:
+        pass
+
+    @abstractmethod
+    def source(self) -> vs.VideoNode:
+        pass
+
+
+class SimpleSource(FileSource):
+    src: List[FileTrim]
+
+    def __init__(self, src: Union[FileTrim, List[FileTrim]]) -> None:
+        self.src = src if isinstance(src, list) else [src]
+
+    def get_audio(self) -> List[FileTrim]:
+        return self.src
+
+    def source(self) -> vs.VideoNode:
+        return depth(core.std.Splice([s.apply_trim(self._open(s.path)) for s in self.src]), 16)
+
+
+class DehardsubFileFinder(FileSource):
     config: Config
 
     def __init__(self, config: Config) -> None:
@@ -57,14 +102,17 @@ class DehardsubFileFinder(ABC):
     def get_ref(self) -> vs.VideoNode:
         pass
 
-    def source(self) -> Tuple[List[vs.VideoNode], vs.VideoNode]:
+    def source(self) -> vs.VideoNode:
+        return self.get_ref()
+
+    def dhs_source(self) -> Tuple[List[vs.VideoNode], vs.VideoNode]:
         wakas: List[vs.VideoNode] = []
         for f in [self.config.format_filename(f) for f in self.get_waka_filenames()]:
             if not os.path.isfile(f):
                 log.warn("Missing a waka!")
                 continue
-            wakas.append(vsutil.depth(core.ffms2.Source(f), 16))
+            wakas.append(self._open(f))
-        ref = vsutil.depth(self.get_ref(), 16)
+        ref = self.get_ref()
         return wakas, ref
 
 
@@ -75,12 +123,24 @@ class FunimationSource(DehardsubFileFinder):
         self.ref_is_funi = False
         super().__init__(*args, **kwargs)
 
+    def get_audio(self) -> List[FileTrim]:
+        if self.ref_is_funi:
+            return [FileTrim(self.get_funi_filename(), (FUNI_INTRO, None))]
+
+        if os.path.isfile(self.config.format_filename(AMAZON_FILENAME_CBR)):
+            return [FileTrim(self.config.format_filename(AMAZON_FILENAME_CBR), None)]
+
+        if os.path.isfile(self.config.format_filename(AMAZON_FILENAME_VBR)):
+            return [FileTrim(self.config.format_filename(AMAZON_FILENAME_VBR), None)]
+
+        raise FileNotFoundError("Failed to find audio that should exist!")
+
     def get_amazon(self) -> vs.VideoNode:
         if not os.path.isfile(self.config.format_filename(AMAZON_FILENAME_CBR)):
             log.warn("Amazon not found, falling back to Funimation")
             raise FileNotFoundError()
         log.success("Found Amazon video")
-        return core.ffms2.Source(self.config.format_filename(AMAZON_FILENAME_CBR))
+        return self._open(self.config.format_filename(AMAZON_FILENAME_CBR))
 
     def get_funi_filename(self) -> str:
         try:
@@ -95,7 +155,7 @@ class FunimationSource(DehardsubFileFinder):
             raise
 
     def get_funi(self) -> vs.VideoNode:
-        return core.ffms2.Source(self.get_funi_filename())[FUNI_INTRO:]
+        return self._open(self.get_funi_filename())[FUNI_INTRO:]
 
     def get_ref(self) -> vs.VideoNode:
         try: