realtime-RIFE-portable/mpv/interpolation.vpy at main · We0M/realtime-RIFE-portable · GitHub

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
import vapoursynth as vs
from vapoursynth import core
import vsmlrt
from vsmlrt import RIFE, RIFEModel, Backend
import os
from fractions import Fraction

# =============== Settings ===============

core.max_cache_size = 8192 #! 8Gb RAM

SC_THRESHOLD   = 0.10 # Recommended range: 0.08-0.16
# RIFE
RIFE_V         = 2 # _implementation
RIFE_FORMAT    = vs.RGBH # vs.RGBS not good for realtime on GPU < RTX3080 if TRT_FP16==False
RIFE_MULTI     = 2
RIFE_MODEL     = 425
# TensorRT
TRT_FP16       = True
TRT_STREAMS    = 2
TRT_CUDA_GRAPH = True
OUTPUT_FORMAT  = 1 # (0 - fp32,1 -fp16)

# ================ Source ================

if 'video_in' in globals():
    # MPV
    clip = video_in
else:
    src = globals().get('src', None)

    if src and os.path.isfile(src):
        clip = core.lsmas.LWLibavSource(src)
    else:
        raise ValueError(
            "Video not found!\n"
            "Example:\n"
            '  vspipe -c y4m rife.vpy - -a src="D:/video.mkv"\n'
            "  mpv --vf=vapoursynth=rife.vpy video.mkv"
        )

# ================= FPS ==================

if 'video_in' in globals():
    fps = Fraction(container_fps).limit_denominator(1000)
    fps_num = fps.numerator
    fps_den = fps.denominator
else:
    fps_num = clip.fps.numerator
    fps_den = clip.fps.denominator

# ================ Matrix ================

is_hd = clip.width > 1024 or clip.height >= 600
est_matrix = 1 if is_hd else 5

def ensure_matrix(clip):
    m = clip.get_frame(0).props.get('_Matrix')
    if m is None or m == 2:
        clip = core.std.SetFrameProps(clip, _Matrix=est_matrix)
    return clip

clip = ensure_matrix(clip)

# ========== Scene Change Detect =========

if clip.format.color_family in [vs.YUV, vs.GRAY]:
    sc_clip = core.std.ShufflePlanes(clip, planes=0, colorfamily=vs.GRAY)
else:
    sc_clip = core.resize.Bilinear(clip, format=vs.GRAY8, matrix=est_matrix)

sc_clip = sc_clip.misc.SCDetect(SC_THRESHOLD)

# ============== AddPadding ==============

# RIFE.TRT requires sizes that are multiples of 32(sometimes 64 or 128).
if RIFE_MODEL == 4251:
    mod = 128
else:
    mod = 32

pad_w = (-clip.width ) % mod
pad_h = (-clip.height) % mod

if pad_w or pad_h:
    clip = core.std.AddBorders(clip, right=pad_w, bottom=pad_h)

# ================ Resize ================

clip = core.resize.Bicubic(clip, format=RIFE_FORMAT)

clip = clip.std.CopyFrameProps(sc_clip, ["_SceneChangePrev", "_SceneChangeNext"])

# ============ Interpolation =============

clip = RIFE(
    clip,
    multi=RIFE_MULTI,
    model=RIFE_MODEL,
    _implementation=RIFE_V,
    backend=Backend.TRT(
        fp16=TRT_FP16,
        use_cuda_graph=TRT_CUDA_GRAPH,
        static_shape=True,
        num_streams=TRT_STREAMS,
        output_format=OUTPUT_FORMAT,
    ),
)

clip = core.std.AssumeFPS(clip, fpsnum=(fps_num * RIFE_MULTI), fpsden=fps_den)

# ============= Crop Padding =============

if pad_w or pad_h:
    clip = core.std.Crop(clip, right=pad_w, bottom=pad_h)

# ================ Output ================

clip = core.resize.Bicubic(clip, format=vs.YUV420P10, matrix=est_matrix)
clip.set_output()