Skip to content

Convert Rife Models to NCNN format

TNTwise edited this page May 3, 2024 · 1 revision

Setup

  • First, we need to clone the Rife github. git clone https://github.com/hzwer/Practical-RIFE.git && cd Practical-RIFE
  • pip install -r requirements.txt
  • We will be using the Rife-v4.6 model, but I have added additional information to convert newer models, which you can download here: https://drive.google.com/file/d/1EAbsfY7mjnXNa6RAsATj2ImAEqmHTjbE/view
  • move train_log into the practical rife directory

Convert To ONNX

  • Edit the file RIFE_HDV3.py in the train_log directory with any text editor
  • delete the call to self.flownet in the inference function.
  • Add this code in the inference function before return merged[3]:
     test_input = torch.rand(1, 3, 256, 256)
     timestep = torch.Tensor([0.5])
     torch.onnx.export(
self.flownet,
(test_input,test_input, torch.Tensor([0.5])),
"rife.onnx",
verbose=False,
opset_version=11,
input_names=["in0", "in1", "in2"],
output_names=["out0"],)
  • if you are exporting with pnnx:
test_input = torch.rand(1, 3, 256, 256)
        timestep = torch.Tensor([1])
        mod = torch.jit.trace(self.flownet,(test_input,test_input,torch.Tensor([1]))) 
        mod.save('rife.pt')

  • Edit the IFnet file:
  • When calling forward in IFnet, remove the x and replace with img0,img1
  • Delete
if training == False:
    channel = x.shape[1] // 2
    img0 = x[:, :channel]
    img1 = x[:, channel:]
  • This redefines the image variable, and we do not want this.
  • copy x = torch.cat((img0,img1),1) to top of forward function, this recalculates the x that we removed from the function call
  • Replace timestep
timestep = (x[:, :1].clone() * 0 + 1) * timestep
timestep = timestep.float()
  • Replace warp

    warped_img0 = warp(img0, flow[:, :2])
    warped_img1 = warp(img1, flow[:, 2:4])
    
    • to
    warped_img0 = img0**flow[:,:3]
    warped_img1 = img1**flow[:,1:4]
    

    If applicable: (<v4.6)

    wf0 = warp(f0, flow[:, :2])
    wf1 = warp(f1, flow[:, 2:4])
    
    • to
    wf0 = f0**flow[:, 1:2]
    wf1 = f1**flow[:, 1:2]
    

    if using pnnx to convert:

    wf0 = f0**flow[:, 1:2]
    wf1 = f1**flow[:, 2:3]
    
  • at the return statement at the bottom, delete the line, and write return merged[3]

return flow_list, mask_list[3], merged

to

return merged[3]

If you are exporting with onnx, include this in the forward fucntion of the ifblock (excluding class ifblock and def forward lines)

class IFBlock(nn.Module):

    def forward(self, x, flow=None, scale=1):
        try:
            scale = scale.item()
        except Exception as e:
            pass
wget https://raw.githubusercontent.com/TNTwise/Rife-NCNN-Model-Comparisons/c053eaf9b51fa07467954d4d8ed1cf752b1fd68b/0.png && wget https://raw.githubusercontent.com/TNTwise/Rife-NCNN-Model-Comparisons/c053eaf9b51fa07467954d4d8ed1cf752b1fd68b/2.png
  • Run code with test images provided python3 inference_img.py --img 0.png 2.png --exp 1

Convert onnx to ncnn

  • Go to convertmodel.com, select the onnx file outputed, and check simplify, optimize and generate fp16 model.
  • Download the bin and param file

Convert to ncnn with pnnx:

./pnnx rife.pt inputshape=[1,3,256,256],[1,3,256,256],[1] fp16=1 optlevel=2

Cropping param file

  • In the param file, fix the cropping by replacing every Pow operator with warp, and replace every corresponding BinaryOp operator to rife.Warp like so:
    BinaryOp                 Pow_189                  2 1 in0_splitncnn_3 335 336 0=6
    
    • to
    rife.Warp                 warp_189                  2 1 in0_splitncnn_3 335 336 0=6
    
  • Do this for every instance of Pow
  • Fix cropping by going from this:
#1,0 / 1,3 = [:, :3]
Crop                     Slice_188                1 1 325_splitncnn_7 335 -23309=1,0 -23310=1,3 -23311=1,0
rife.Warp                 warp_189                  2 1 in0_splitncnn_3 335 336 0=6

#1,1 / 1,1 = [:, 1:4]
Crop                     Slice_194                1 1 325_splitncnn_6 341 -23309=1,1 -23310=1,4 -23311=1,0
rife.Warp                warp_195                  2 1 in1_splitncnn_3 341 342 0=6

To This:

#1,0 / 1,2 = [:, :2]
Crop                     Slice_188                1 1 325_splitncnn_7 335 -23309=1,0 -23310=1,2 -23311=1,0
rife.Warp                 warp_189                  2 1 in0_splitncnn_3 335 336 0=6

# 1,2 / 1,4 = [:, 2:4]
Crop                     Slice_194                1 1 325_splitncnn_6 341 -23309=1,2 -23310=1,4 -23311=1,0
rife.Warp                warp_195                  2 1 in1_splitncnn_3 341 342 0=6
  • If you are converting > 4.6 Crop the other warp function
#1,1 / 1,2 = [:, 1:2]
Crop                     Slice_200                1 1 325_splitncnn_5 347 -23309=1,1 -23310=1,2 -23311=1,0
rife.Warp                 warp_201                  2 1 154_splitncnn_2 347 348 0=6

#1,1 / 1,2 = [:, 1:2]
Crop                     Slice_206                1 1 325_splitncnn_4 353 -23309=1,1 -23310=1,2 -23311=1,0
rife.Warp                 warp_207                  2 1 166_splitncnn_2 353 354 0=6

To This:

#1,0 / 1,2 = [:, :2]
Crop                     Slice_200                1 1 325_splitncnn_5 347 -23309=1,0 -23310=1,2 -23311=1,0
rife.Warp                 warp_201                  2 1 154_splitncnn_2 347 348 0=6

#1,2 / 1,4 = [:, 2:4]
Crop                     Slice_206                1 1 325_splitncnn_4 353 -23309=1,2 -23310=1,4 -23311=1,0
rife.Warp                 warp_207                  2 1 166_splitncnn_2 353 354 0=6

if pnnx:

BinaryOp                 pow_43                   2 1 4 201 202 0=6
Crop                     slice_120                1 1 199 203 -23310=1,4 -23311=1,0 -23309=1,1
BinaryOp                 pow_44                   2 1 12 203 204 0=6
Crop                     slice_121                1 1 198 205 -23310=1,2 -23311=1,0 -23309=1,1
BinaryOp                 pow_45                   2 1 40 205 206 0=6
Crop                     slice_122                1 1 197 207 -23310=1,3 -23311=1,0 -23309=1,2
BinaryOp                 pow_46                   2 1 32 207 208 0=6

to

Crop                     slice_119                1 1 200 201 -23310=1,2 -23311=1,0 -23309=1,0
rife.Warp                warp_43                   2 1 4 201 202 0=6
Crop                     slice_120                1 1 199 203 -23310=1,4 -23311=1,0 -23309=1,2
rife.Warp                warp_44                   2 1 12 203 204 0=6
Crop                     slice_121                1 1 198 205 -23310=1,2 -23311=1,0 -23309=1,2
rife.Warp                warp_45                   2 1 40 205 206 0=6
Crop                     slice_122                1 1 197 207 -23310=1,4 -23311=1,0 -23309=1,2
rife.Warp                warp_46                   2 1 32 207 208 0=6
  • Repeat until every warp is cropped.
  • Rename bin and param files to flownet.bin and flownet.param
  • Put the models in a folder called rife-v4-something, It has to include rife-v4 as rife-ncnn-vulkan chooses what kind of model it is based on the name.
  • Done! you can now run the model with the binary here, nihui's won't work as it is not compiled with the MemoryData operator.
  • To use with REAL Video Enhancer, move the model to:
  • AppImage: ~/.local/share/REAL-Video-Enhancer/models/rife/
  • Flatpak: ~/.var/app/io.github.tntwise.REAL-Video-Enhancer/models/rife/

Resources:

https://github.com/styler00dollar/VSGAN-tensorrt-docker/issues/37
https://github.com/nihui/rife-ncnn-vulkan/issues/50#issuecomment-1229203013
https://github.com/styler00dollar/VSGAN-tensorrt-docker/releases/tag/models
https://zhuanlan.zhihu.com/p/93017149
https://github.com/styler00dollar/VapourSynth-RIFE-ncnn-Vulkan/issues/6