-
-
Notifications
You must be signed in to change notification settings - Fork 6
Convert Rife Models to NCNN format
TNTwise edited this page May 3, 2024
·
1 revision
- 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
- 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
- Go to convertmodel.com, select the onnx file outputed, and check simplify, optimize and generate fp16 model.
- Download the bin and param file
./pnnx rife.pt inputshape=[1,3,256,256],[1,3,256,256],[1] fp16=1 optlevel=2
- 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
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/
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