# Imports
import vapoursynth as vs
# getting Vapoursynth core
import site
import sys
import os
core = vs.core
# Import scripts folder
scriptPath = 'F:/Hybrid/64bit/vsscripts'
sys.path.insert(0, os.path.abspath(scriptPath))
os.environ["CUDA_MODULE_LOADING"] = "LAZY"
# loading plugins
core.std.LoadPlugin(path="F:/Hybrid/64bit/vs-mlrt/vstrt.dll")
core.std.LoadPlugin(path="F:/Hybrid/64bit/vsfilters/Support/fmtconv.dll")
core.std.LoadPlugin(path="F:/Hybrid/64bit/vsfilters/FrameFilter/RIFE/librife_r9_mod_v6.dll")
core.std.LoadPlugin(path="F:/Hybrid/64bit/vsfilters/MiscFilter/MiscFilters/MiscFilters.dll")
core.std.LoadPlugin(path="F:/Hybrid/64bit/vsfilters/DeinterlaceFilter/TIVTC/libtivtc.dll")
core.std.LoadPlugin(path="F:/Hybrid/64bit/vsfilters/SourceFilter/LSmashSource/LSMASHSource.dll")
# Import scripts
from importlib.machinery import SourceFileLoader
vsmlrt = SourceFileLoader('vsmlrt', 'F:/Hybrid/64bit/vs-mlrt/vsmlrt.py').load_module()
import validate
# Source: 'C:\Users\Selur\Desktop\VH.avi'
# Current color space: YUV444P10, bit depth: 10, resolution: 760x488, frame rate: 29.97fps, scanorder: telecine, yuv luminance scale: limited, matrix: 470bg
# Loading C:\Users\Selur\Desktop\VH.avi using LWLibavSource
clip = core.lsmas.LWLibavSource(source="C:/Users/Selur/Desktop/VH.avi", format="YUV444P10", stream_index=0, cache=0, prefer_hw=0)
frame = clip.get_frame(0)
# Setting detected color matrix (470bg).
clip = core.std.SetFrameProps(clip=clip, _Matrix=5)
# setting color transfer (170), if it is not set.
if validate.transferIsInvalid(clip):
  clip = core.std.SetFrameProps(clip=clip, _Transfer=6)
# setting color primaries info (to 470), if it is not set.
if validate.primariesIsInvalid(clip):
  clip = core.std.SetFrameProps(clip=clip, _Primaries=5)
# setting color range to TV (limited) range.
clip = core.std.SetFrameProps(clip=clip, _ColorRange=1)
# making sure frame rate is set to 29.97fps
clip = core.std.AssumeFPS(clip=clip, fpsnum=30000, fpsden=1001)
# making sure the detected scan type is set (detected: telecine)
clip = core.std.SetFrameProps(clip=clip, _FieldBased=2) # tff
clip = core.tivtc.TFM(clip=clip)
# Making sure content is preceived as frame based
clip = core.std.SetFrameProps(clip=clip, _FieldBased=0) # progressive
# adjusting color space from YUV444P10 to RGBS for vsFillDuplicateFrames
clip = core.resize.Bicubic(clip=clip, format=vs.RGBS, matrix_in_s="470bg", range_s="limited")
from FillDuplicateFrames import FillDuplicateFrames
fdf = FillDuplicateFrames(clip=clip, method="RIFE", thresh=0.020000)
clip = fdf.out
# changing range from limited to full range for vsCodeFormer
clip = core.resize.Bicubic(clip, range_in_s="limited", range_s="full")
# adjusting color space from RGBS to RGB24 for vsCodeFormer
clip = core.resize.Bicubic(clip=clip, format=vs.RGB24, range_s="full", dither_type="error_diffusion")
# Blind Face Restoration using CodeFormer
# Using background resizer FeMaSR(2x)
clipBG = clip
from vsfemasr import femasr as FeMaSR
# adjusting color space from RGB24 to RGBH for vsFeMaSR
clipBG = core.resize.Bicubic(clip=clipBG, format=vs.RGBH, range_s="full")
# Step 1: current: 760x488, target: 1520x976
# resizing using FeMaSR
clipBG = FeMaSR(clip=clipBG, device_index=0) # 1520x976
# resizing 1520x976 to 1520x976
# adjusting resizing
clipBG = core.resize.Bicubic(clip=clipBG, format=vs.RGBS, range_s="full")
clipBG = core.fmtc.resample(clip=clipBG, w=1520, h=976, kernel="spline64", interlaced=False, interlacedd=False)
clipBG = core.resize.Bicubic(clip=clipBG, format=vs.RGB24, range_s="full", dither_type="error_diffusion")
from vscodeformer import codeformer as CodeFormer
clip = CodeFormer(clip=clip, upscale=2, detector=1, weight=1.000, bg_clip=clipBG) # 1520x976
# changing range from full to limited range for vsCodeFormer
clip = core.resize.Bicubic(clip, range_in_s="full", range_s="limited")
from vsmlrt import Backend
# adjusting color space from RGB24 to RGBH for vsSCUNetmlrt
clip = core.resize.Bicubic(clip=clip, format=vs.RGBH, range_s="limited")
# denosing using SCUNet (mlrt)
clip = vsmlrt.SCUNet(clip, model=4, overlap=16, backend=Backend.TRT(fp16=True,device_id=0,verbose=True,use_cuda_graph=False, num_streams=3,builder_optimization_level=3,engine_folder="J:/tmp"))
# Resizing using 10 - bicubic spline
clip = core.resize.Bicubic(clip=clip, format=vs.RGBS, range_s="limited")
clip = core.fmtc.resample(clip=clip, kernel="spline16", w=1280, h=822, interlaced=False, interlacedd=False) # resolution 1280x822
# adjusting output color from: RGBS to YUV420P10 for NVEncModel
clip = core.resize.Bicubic(clip=clip, format=vs.YUV420P10, matrix_s="470bg", range_s="limited", dither_type="error_diffusion")
# set output frame rate to 29.97fps (progressive)
clip = core.std.AssumeFPS(clip=clip, fpsnum=30000, fpsden=1001)
# output
clip.set_output()