diff --git a/export_onnx.py b/export_onnx.py index 97f2948..96b77fc 100644 --- a/export_onnx.py +++ b/export_onnx.py @@ -14,10 +14,10 @@ from onnxsim import simplify MODEL_TYPE = pathlib.Path(__file__).parent.name CKPT_NAME_VER_DICT = { "rife40.pth": "4.0", - "rife41.pth": "4.0", - "rife42.pth": "4.2", - "rife43.pth": "4.3", - "rife44.pth": "4.3", + "rife41.pth": "4.0", + "rife42.pth": "4.2", + "rife43.pth": "4.3", + "rife44.pth": "4.3", "rife45.pth": "4.5", "rife46.pth": "4.6", "rife47.pth": "4.7", @@ -106,7 +106,7 @@ def export_onnx(ckpt_name, ensemble, scale_factor): 'img1': {2: 'height', 3: 'width'}, 'output': {2: 'height', 3: 'width'}, } - + with torch.no_grad(): # Disable gradients for efficiency torch.onnx.export(interpolation_model, (img0, img1, timestep), # ensemble + scale_factor set in forward fn @@ -124,7 +124,7 @@ def export_onnx(ckpt_name, ensemble, scale_factor): onnx_model = onnx.load(onnx_save_path) onnx.checker.check_model(onnx_model) # Perform a validity check print("ONNX model validation successful!") - + # print(onnx.helper.printable_graph(onnx_model.graph)) sim_model_path = os.path.join("models", f"{ckpt_name.split('.')[0]}_ensemble_{ensemble}_scale_{scale_factor}_sim.onnx") diff --git a/rife_arch.py b/rife_arch.py index 5b1d722..45a2f10 100644 --- a/rife_arch.py +++ b/rife_arch.py @@ -532,6 +532,6 @@ class IFNet(nn.Module): tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) res = tmp[:, :3] * 2 - 1 merged[3] = torch.clamp(merged[3] + res, 0, 1) - + result = merged[3][:, :, :h, :w] return result \ No newline at end of file diff --git a/vfi_utilities.py b/vfi_utilities.py index f9ec616..ae98f74 100644 --- a/vfi_utilities.py +++ b/vfi_utilities.py @@ -7,7 +7,7 @@ import einops from comfy.model_management import soft_empty_cache, get_torch_device import numpy as np from comfy.utils import ProgressBar -from colored import Fore, Back, Style +from colored import Fore, Back, Style DEVICE = get_torch_device() @@ -40,7 +40,7 @@ def generate_frames_rife( multiplier, return_middle_frame_function ): - + output_frames = torch.zeros(multiplier*frames.shape[0], *frames.shape[1:], device="cpu") out_len = 0 @@ -70,7 +70,7 @@ def generate_frames_rife( logger("Clearing cache...") pbar.update(1) - + # Append final frame output_frames[out_len] = frames[-1:] @@ -80,6 +80,6 @@ def generate_frames_rife( # clear cache for courtesy soft_empty_cache() logger("Final clearing cache done ...") -# +# res = output_frames[:out_len] return res \ No newline at end of file