[Bugfix] Fix RIFE device selection for CPU-transported videos (#2876)

Signed-off-by: david6666666 <530634352@qq.com>
This commit is contained in:
WeiQing Chen authored and GitHub committed 2026-04-17 19:19:34 +08:00
1 parent c0ccbb872f
commit b7f239810d
2 files changed
+9 -5

No files matched your search

@@ -72,7 +72,7 @@ def test_frame_interpolator_runs_actual_torch_tensor_path(monkeypatch):
assert torch.isfinite(output_video).all()
def test_frame_interpolator_prefers_input_tensor_device(monkeypatch):
def test_frame_interpolator_uses_platform_device_when_tensor_is_cpu(monkeypatch):
chosen_devices = []
model = rife_interpolator.Model().eval()
@@ -83,10 +83,11 @@ def test_frame_interpolator_prefers_input_tensor_device(monkeypatch):
interpolator = rife_interpolator.FrameInterpolator()
monkeypatch.setattr(interpolator, "_ensure_model_loaded", _fake_ensure_model_loaded)
monkeypatch.setattr(model.flownet, "to", lambda device: model.flownet)
monkeypatch.setattr(rife_interpolator, "_select_torch_device", lambda: torch.device("cuda"))
video = torch.zeros(1, 3, 2, 32, 32)
output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0)
assert chosen_devices == [video.device]
assert chosen_devices == [torch.device("cuda")]
assert multiplier == 2
assert output_video.shape == (1, 3, 3, 32, 32)