diff --git a/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py index a2d80d4b9..e5becaa9d 100644 --- a/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py +++ b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py @@ -243,7 +243,7 @@ def test_stage_attaches_the_branches_to_each_materialized_lazy_transformer(monke assert second.targets[-1] is not loaded_first -def _denoise(monkeypatch, steps, velocity=0.5, transformer=None, curve=None): +def _denoise(monkeypatch, steps, velocity=0.5, transformer=None, curve=None, fastvideo_args=None): from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import denoising @@ -260,7 +260,7 @@ def _denoise(monkeypatch, steps, velocity=0.5, transformer=None, curve=None): transformer = transformer if transformer is not None else _ConstantVelocityTransformer(velocity) curve = curve if curve is not None else _RecordingCurve() stage = denoising.Regenerate2KDenoisingStage(transformer, curve, None) - return stage.forward(batch, None), transformer, curve + return stage.forward(batch, fastvideo_args), transformer, curve def test_euler_steps_integrate_the_constant_velocity_over_sigma_zero_to_start(monkeypatch): @@ -295,3 +295,46 @@ def test_registry_resolves_the_release_to_the_regenerate_pipeline(tmp_path): assert get_model_info(path).pipeline_cls.__name__ == "OpenH3Regenerate2KPipeline" preset, family = get_preset_selection(hub_id) assert (preset, family) == ("open_h3_regenerate_2k", "open_h3_regenerate_2k") + + +class _OffloadedTransformer(_ConstantVelocityTransformer): + + def __init__(self, velocity, fail=False): + super().__init__(velocity) + self.fail = fail + self.moves = [] + + def to(self, device): + self.moves.append(str(device)) + return self + + def __call__(self, **kwargs): + if self.fail: + raise RuntimeError("forward failed") + return super().__call__(**kwargs) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_full_cpu_offload_moves_the_transformer_back_to_cpu_after_the_refine(monkeypatch, fail): + from types import SimpleNamespace + + transformer = _OffloadedTransformer(0.5, fail=fail) + args = SimpleNamespace(dit_cpu_offload=True, dit_layerwise_offload=False, use_fsdp_inference=False) + if fail: + with pytest.raises(RuntimeError, match="forward failed"): + _denoise(monkeypatch, steps=2, transformer=transformer, fastvideo_args=args) + else: + _denoise(monkeypatch, steps=2, transformer=transformer, fastvideo_args=args) + + assert transformer.moves == ["cpu", "cpu"] # to the stage device (CPU in this test), then back to the host + + +@pytest.mark.parametrize("variable, value", [("FASTVIDEO_H3_ADALN_CACHE", True), + ("FASTVIDEO_H3_ADALN_TABLE", "/tmp/adaln_tables.pt")]) +def test_pipeline_refuses_an_adaln_cache_that_would_skip_the_sharp_branches(variable, value): + from fastvideo import envs + from fastvideo.pipelines.basic.open_h3_regenerate_2k.open_h3_regenerate_2k_pipeline import refuse_adaln_cache + + refuse_adaln_cache() + with getattr(envs, variable).override(value), pytest.raises(ValueError, match="Sharp curve branches"): + refuse_adaln_cache()