diff --git a/Code/OS/AiProcess.cs b/Code/OS/AiProcess.cs index c2d71b3..669abdc 100644 --- a/Code/OS/AiProcess.cs +++ b/Code/OS/AiProcess.cs @@ -125,11 +125,12 @@ namespace Flowframes { //bool parallel = false; string outPath = Path.Combine(inPath.GetParentDir(), outDir); + string uhdStr = await InterpolateUtils.UseUHD() ? "--UHD" : ""; string wthreads = $"--wthreads {2 * (int)interpFactor}"; string rbuffer = $"--rbuffer {Config.GetInt("rifeCudaBufferSize", 200)}"; - string scale = $"--scale {Config.GetFloat("rifeCudaScale", 1.0f).ToStringDot()}"; + //string scale = $"--scale {Config.GetFloat("rifeCudaScale", 1.0f).ToStringDot()}"; string prec = Config.GetBool("rifeCudaFp16") ? "--fp16" : ""; - string args = $" --input {inPath.Wrap()} --output {outDir} --model {mdl} --exp {(int)Math.Log(interpFactor, 2)} {scale} {wthreads} {rbuffer} {prec}"; + string args = $" --input {inPath.Wrap()} --output {outDir} --model {mdl} --exp {(int)Math.Log(interpFactor, 2)} {uhdStr} {wthreads} {rbuffer} {prec}"; // if (parallel) args = $" --input {inPath.Wrap()} --output {outPath} --model {mdl} --factor {interpFactor}"; // if (parallel) script = "rife-parallel.py"; @@ -138,7 +139,7 @@ namespace Flowframes SetProgressCheck(Path.Combine(Interpolate.current.tempFolder, outDir), interpFactor); rifePy.StartInfo.Arguments = $"{OSUtils.GetCmdArg()} cd /D {PkgUtils.GetPkgFolder(Packages.rifeCuda).Wrap()} & " + $"set CUDA_VISIBLE_DEVICES={Config.Get("torchGpus")} & {Python.GetPyCmd()} {script} {args}"; - Logger.Log($"Running RIFE (CUDA)...".TrimWhitespaces(), false); + Logger.Log($"Running RIFE (CUDA){(await InterpolateUtils.UseUHD() ? " (UHD Mode)" : "")}...", false); Logger.Log("cmd.exe " + rifePy.StartInfo.Arguments, true); if (!OSUtils.ShowHiddenCmd()) diff --git a/Pkgs/rife-cuda/model/IFNet_HD.py b/Pkgs/rife-cuda/model/IFNet_HD.py index 6975679..fe315b2 100644 --- a/Pkgs/rife-cuda/model/IFNet_HD.py +++ b/Pkgs/rife-cuda/model/IFNet_HD.py @@ -91,9 +91,12 @@ class IFNet(nn.Module): self.block2 = IFBlock(8, scale=2, c=96) self.block3 = IFBlock(8, scale=1, c=48) - def forward(self, x, scale=1.0): - x = F.interpolate(x, scale_factor=0.5 * scale, mode="bilinear", - align_corners=False) + def forward(self, x, UHD=False): + if UHD: + x = F.interpolate(x, scale_factor=0.25, mode="bilinear", align_corners=False) + else: + x = F.interpolate(x, scale_factor=0.5, mode="bilinear", + align_corners=False) flow0 = self.block0(x) F1 = flow0 warped_img0 = warp(x[:, :3], F1) @@ -108,8 +111,6 @@ class IFNet(nn.Module): warped_img1 = warp(x[:, 3:], -F3) flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3), 1)) F4 = (flow0 + flow1 + flow2 + flow3) - F4 = F.interpolate(F4, scale_factor=1 / scale, mode="bilinear", - align_corners=False) / scale return F4, [F1, F2, F3, F4] if __name__ == '__main__': diff --git a/Pkgs/rife-cuda/model/IFNet_HDv2.py b/Pkgs/rife-cuda/model/IFNet_HDv2.py index c71886c..f9b18cf 100644 --- a/Pkgs/rife-cuda/model/IFNet_HDv2.py +++ b/Pkgs/rife-cuda/model/IFNet_HDv2.py @@ -39,19 +39,18 @@ class IFBlock(nn.Module): ) self.conv1 = nn.ConvTranspose2d(2*c, 4, 4, 2, 1) - def forward(self, x, scale=1.0): - infer_scale = self.scale / scale - if infer_scale != 1.0: - x = F.interpolate(x, scale_factor=1. / infer_scale, mode="bilinear", + def forward(self, x): + if self.scale != 1: + x = F.interpolate(x, scale_factor=1. / self.scale, mode="bilinear", align_corners=False) x = self.conv0(x) x = self.convblock(x) x = self.conv1(x) flow = x - if infer_scale != 1.0: - flow = F.interpolate(flow, scale_factor=infer_scale, mode="bilinear", - align_corners=False) - return flow / scale + if self.scale != 1: + flow = F.interpolate(flow, scale_factor=self.scale, mode="bilinear", + align_corners=False) + return flow class IFNet(nn.Module): @@ -62,23 +61,25 @@ class IFNet(nn.Module): self.block2 = IFBlock(10, scale=2, c=96) self.block3 = IFBlock(10, scale=1, c=48) - def forward(self, x, scale=1.0): - flow0 = self.block0(x, scale) + def forward(self, x, UHD=False): + if UHD: + x = F.interpolate(x, scale_factor=0.5, mode="bilinear", align_corners=False) + flow0 = self.block0(x) F1 = flow0 F1_large = F.interpolate(F1, scale_factor=2.0, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 2.0 warped_img0 = warp(x[:, :3], F1_large[:, :2]) warped_img1 = warp(x[:, 3:], F1_large[:, 2:4]) - flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1_large), 1), scale) + flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1_large), 1)) F2 = (flow0 + flow1) F2_large = F.interpolate(F2, scale_factor=2.0, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 2.0 warped_img0 = warp(x[:, :3], F2_large[:, :2]) warped_img1 = warp(x[:, 3:], F2_large[:, 2:4]) - flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2_large), 1), scale) + flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2_large), 1)) F3 = (flow0 + flow1 + flow2) F3_large = F.interpolate(F3, scale_factor=2.0, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 2.0 warped_img0 = warp(x[:, :3], F3_large[:, :2]) warped_img1 = warp(x[:, 3:], F3_large[:, 2:4]) - flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3_large), 1), scale) + flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3_large), 1)) F4 = (flow0 + flow1 + flow2 + flow3) return F4, [F1, F2, F3, F4] diff --git a/Pkgs/rife-cuda/model/RIFE_HD.py b/Pkgs/rife-cuda/model/RIFE_HD.py index 461dc1b..b96576f 100644 --- a/Pkgs/rife-cuda/model/RIFE_HD.py +++ b/Pkgs/rife-cuda/model/RIFE_HD.py @@ -188,9 +188,11 @@ class Model: torch.save(self.contextnet.state_dict(), '{}/contextnet.pkl'.format(path)) torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path)) - def predict(self, imgs, flow, training=True, flow_gt=None): + def predict(self, imgs, flow, training=True, flow_gt=None, UHD=False): img0 = imgs[:, :3] img1 = imgs[:, 3:] + if UHD: + flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0 c0 = self.contextnet(img0, flow) c1 = self.contextnet(img1, -flow) flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear", @@ -207,10 +209,10 @@ class Model: else: return pred - def inference(self, img0, img1, scale=1.0): + def inference(self, img0, img1, UHD=False): imgs = torch.cat((img0, img1), 1) - flow, _ = self.flownet(imgs, scale) - return self.predict(imgs, flow, training=False) + flow, _ = self.flownet(imgs, UHD) + return self.predict(imgs, flow, training=False, UHD=UHD) def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): for param_group in self.optimG.param_groups: diff --git a/Pkgs/rife-cuda/model/RIFE_HDv2.py b/Pkgs/rife-cuda/model/RIFE_HDv2.py index 2de119d..9f19ae2 100644 --- a/Pkgs/rife-cuda/model/RIFE_HDv2.py +++ b/Pkgs/rife-cuda/model/RIFE_HDv2.py @@ -173,9 +173,11 @@ class Model: torch.save(self.contextnet.state_dict(), '{}/contextnet.pkl'.format(path)) torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path)) - def predict(self, imgs, flow, training=True, flow_gt=None): + def predict(self, imgs, flow, training=True, flow_gt=None, UHD=False): img0 = imgs[:, :3] img1 = imgs[:, 3:] + if UHD: + flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0 c0 = self.contextnet(img0, flow[:, :2]) c1 = self.contextnet(img1, flow[:, 2:4]) flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear", @@ -192,10 +194,10 @@ class Model: else: return pred - def inference(self, img0, img1, scale=1.0): + def inference(self, img0, img1, UHD=False): imgs = torch.cat((img0, img1), 1) - flow, _ = self.flownet(imgs, scale) - return self.predict(imgs, flow, training=False) + flow, _ = self.flownet(imgs, UHD) + return self.predict(imgs, flow, training=False, UHD=UHD) def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): for param_group in self.optimG.param_groups: diff --git a/Pkgs/rife-cuda/rife.py b/Pkgs/rife-cuda/rife.py index dd531f6..c2437db 100644 --- a/Pkgs/rife-cuda/rife.py +++ b/Pkgs/rife-cuda/rife.py @@ -107,7 +107,7 @@ def build_read_buffer(user_args, read_buffer, videogen): def make_inference(I0, I1, exp): global model - middle = model.inference(I0, I1, args.scale) + middle = model.inference(I0, I1, args.UHD) if exp == 1: return [middle] first_half = make_inference(I0, middle, exp=exp - 1) @@ -120,10 +120,13 @@ def pad_image(img): else: return F.pad(img, padding) -print(f"Scale: {args.scale}") -tmp = max(32, int(32 / args.scale)) -ph = ((h - 1) // tmp + 1) * tmp -pw = ((w - 1) // tmp + 1) * tmp +if args.UHD: + print("UHD mode enabled.") + ph = ((h - 1) // 64 + 1) * 64 + pw = ((w - 1) // 64 + 1) * 64 +else: + ph = ((h - 1) // 32 + 1) * 32 + pw = ((w - 1) // 32 + 1) * 32 padding = (0, pw - w, 0, ph - h) write_buffer = Queue(maxsize=args.rbuffer)