From 9eefdbddfa5a8ced418ce53b3f0896af1a2d663c Mon Sep 17 00:00:00 2001 From: huchukato Date: Wed, 5 Aug 2026 11:01:43 +0200 Subject: [PATCH] Filter training-only keys (teacher.*, caltime.*) from checkpoints rife426.pth and possibly other new checkpoints contain extra modules used during training (teacher network, caltime scheduler) that are not part of the IFNet inference architecture. These cause load_state_dict to fail with "Unexpected key(s)" errors. Now export_onnx.py filters the state_dict to keep only keys that match the model architecture before loading, with strict=True to catch any real missing keys. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- export_onnx.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/export_onnx.py b/export_onnx.py index 87426f3..0442e60 100644 --- a/export_onnx.py +++ b/export_onnx.py @@ -120,7 +120,19 @@ def export_onnx(ckpt_name, ensemble, scale_factor): model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) arch_ver = CKPT_NAME_VER_DICT[ckpt_name] interpolation_model = IFNet(arch_ver=arch_ver) - interpolation_model.load_state_dict(torch.load(model_path)) + # Some checkpoints (e.g. rife426.pth) contain extra training-only modules + # (teacher.*, caltime.*) that are not part of the inference architecture. + # Filter them out to avoid load_state_dict errors. + raw_state_dict = torch.load(model_path, weights_only=False) + model_keys = set(interpolation_model.state_dict().keys()) + filtered_state_dict = { + k: v for k, v in raw_state_dict.items() + if k in model_keys + } + skipped = set(raw_state_dict.keys()) - model_keys + if skipped: + print(f"Filtered out {len(skipped)} training-only keys: {sorted(skipped)[:5]}...") + interpolation_model.load_state_dict(filtered_state_dict, strict=True) interpolation_model.eval().to(TORCH_DEVICE) # # dummy data