diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index 26e40a4761..213fe12f38 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -614,8 +614,7 @@ def _timed_video( threshold = cache_threshold, mode = auto_cache_mode(family), family = family, - quality = normalize_cache_quality(cache_quality) - or auto_cache_quality(family), + quality = normalize_cache_quality(cache_quality) or auto_cache_quality(family), expert = expert, logger = logger, ) diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index aed5f63cec..2c42a40385 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -319,6 +319,7 @@ def _magcache_ratio_key(family: Optional[str], expert: Optional[str]) -> str: return fam return f"{fam}::{exp}" + # Families whose AUTO step-cache decision engages MagCache instead of FBCache. On # HunyuanVideo-1.5 FBCache free-runs (no skip cap, no error budget) and derails the # trajectory (LPIPS 0.54 + a luma shift at its default threshold), while MagCache holds @@ -667,9 +668,7 @@ def apply_step_cache( # the schedule, so the expert runs ~len(ratios) * steps / 50 forwards. num_steps = int(steps) if len(ratios) != _MAGCACHE_CALIBRATION_STEPS: - num_steps = max( - 1, round(len(ratios) * int(steps) / _MAGCACHE_CALIBRATION_STEPS) - ) + num_steps = max(1, round(len(ratios) * int(steps) / _MAGCACHE_CALIBRATION_STEPS)) config: Any = MagCacheConfig( threshold = thr, max_skip_steps = mag_skip, diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 6fae1ef19b..0597a4461b 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -631,8 +631,7 @@ def test_magcache_ratio_key_primary_and_expert(): assert _magcache_ratio_key("wan2.2-t2v-a14b", None) == "wan2.2-t2v-a14b" assert _magcache_ratio_key("wan2.2-t2v-a14b", "transformer") == "wan2.2-t2v-a14b" assert ( - _magcache_ratio_key("Wan2.2-T2V-A14B", "transformer_2") - == "wan2.2-t2v-a14b::transformer_2" + _magcache_ratio_key("Wan2.2-T2V-A14B", "transformer_2") == "wan2.2-t2v-a14b::transformer_2" ) @@ -643,12 +642,13 @@ def test_magcache_expert_resolves_its_own_curve(monkeypatch): primary_curve = tuple([1.0] * 15) expert_curve = tuple([0.99] * 35) monkeypatch.setitem(dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe", primary_curve) - monkeypatch.setitem( - dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe::transformer_2", expert_curve - ) + monkeypatch.setitem(dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe::transformer_2", expert_curve) t = _MixinTransformer() engaged = apply_step_cache( - _pipe(t), mode = "magcache", family = "fam-moe", steps = 50, + _pipe(t), + mode = "magcache", + family = "fam-moe", + steps = 50, expert = "transformer_2", ) assert engaged == TC_MAGCACHE @@ -664,20 +664,24 @@ def test_magcache_expert_subcurve_scales_step_count(monkeypatch): from core.inference import diffusion_cache as dc_mod expert_curve = tuple([0.99] * 35) - monkeypatch.setitem( - dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe::transformer_2", expert_curve - ) + monkeypatch.setitem(dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe::transformer_2", expert_curve) t = _MixinTransformer() apply_step_cache( - _pipe(t), mode = "magcache", family = "fam-moe", steps = 30, + _pipe(t), + mode = "magcache", + family = "fam-moe", + steps = 30, expert = "transformer_2", ) assert t.enabled_with.num_inference_steps == round(35 * 30 / _MAGCACHE_CALIBRATION_STEPS) # At the calibration step count itself the sub-curve maps 1:1. t2 = _MixinTransformer() apply_step_cache( - _pipe(t2), mode = "magcache", family = "fam-moe", - steps = _MAGCACHE_CALIBRATION_STEPS, expert = "transformer_2", + _pipe(t2), + mode = "magcache", + family = "fam-moe", + steps = _MAGCACHE_CALIBRATION_STEPS, + expert = "transformer_2", ) assert t2.enabled_with.num_inference_steps == 35 @@ -688,7 +692,10 @@ def test_magcache_full_curve_keeps_requested_steps(monkeypatch): _stub_diffusers_with_magcache(monkeypatch) t = _MixinTransformer() apply_step_cache( - _pipe(t), mode = "magcache", family = "wan2.2-ti2v-5b", steps = 30, + _pipe(t), + mode = "magcache", + family = "wan2.2-ti2v-5b", + steps = 30, expert = "transformer", ) assert t.enabled_with.num_inference_steps == 30 @@ -705,7 +712,10 @@ def test_magcache_expert_without_curve_runs_uncached(monkeypatch): t = _MixinTransformer() assert ( apply_step_cache( - _pipe(t), mode = "magcache", family = "fam-moe", steps = 50, + _pipe(t), + mode = "magcache", + family = "fam-moe", + steps = 50, expert = "transformer_2", ) is None @@ -718,12 +728,13 @@ def test_toggle_threads_expert_through(monkeypatch): from core.inference import diffusion_cache as dc_mod expert_curve = tuple([0.98] * 35) - monkeypatch.setitem( - dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe::transformer_2", expert_curve - ) + monkeypatch.setitem(dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe::transformer_2", expert_curve) t = _ToggleTransformer() mode = maybe_toggle_step_cache( - _pipe(t), steps = 50, mode = TC_MAGCACHE, family = "fam-moe", + _pipe(t), + steps = 50, + mode = TC_MAGCACHE, + family = "fam-moe", expert = "transformer_2", ) assert mode == TC_MAGCACHE diff --git a/studio/backend/tests/test_diffusion_cfg_parallel.py b/studio/backend/tests/test_diffusion_cfg_parallel.py index d41406f5bd..8fe2f7cbf7 100644 --- a/studio/backend/tests/test_diffusion_cfg_parallel.py +++ b/studio/backend/tests/test_diffusion_cfg_parallel.py @@ -475,9 +475,7 @@ def test_install_failure_after_cudnn_patch_restores_it(monkeypatch): "_install_threadsafe_cudnn_attention", lambda logger = None: (calls.append("install"), True)[1], ) - monkeypatch.setattr( - cp, "_restore_threadsafe_cudnn_attention", lambda: calls.append("restore") - ) + monkeypatch.setattr(cp, "_restore_threadsafe_cudnn_attention", lambda: calls.append("restore")) class _LoadableDiT(_FakeDiT): @classmethod @@ -492,9 +490,7 @@ def test_install_failure_after_cudnn_patch_restores_it(monkeypatch): pipe = _CtxPipe(_LoadableDiT()) pipe.guider = None # raises AFTER the patch install - proxy, reason = _gate( - monkeypatch, pipe, _fam(), speed_active = False, attention_backend = None - ) + proxy, reason = _gate(monkeypatch, pipe, _fam(), speed_active = False, attention_backend = None) assert proxy is None and reason == "replica install failed" assert calls == ["install", "restore"] diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 7bb8701bcb..301eff4d27 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -603,8 +603,7 @@ def test_select_te_auto_resolves_dense_for_wan_a14b_but_not_wan_5b(monkeypatch): assert select_te_quant_scheme(_target(), "auto", family = "wan2.2-t2v-a14b") is None assert select_te_quant_scheme(_target(), "auto", family = "Wan2.2-T2V-A14B") is None assert ( - select_te_quant_scheme(_target(), "auto", family = "wan2.2-ti2v-5b") - == TE_QUANT_FP8_DYNAMIC + select_te_quant_scheme(_target(), "auto", family = "wan2.2-ti2v-5b") == TE_QUANT_FP8_DYNAMIC ) # The auto-dense table steers only the DEFAULT; an explicit request stays verbatim. assert ( @@ -642,10 +641,7 @@ def test_quantize_explicit_fp8_dynamic_refused_for_ltx2(monkeypatch): calls: list = [] monkeypatch.setattr(dp, "_cast_fp8_dynamic", lambda enc, tgt: calls.append(enc)) pipe = types.SimpleNamespace(text_encoder = object()) - assert ( - quantize_text_encoders(pipe, _target(), mode = "fp8_dynamic", family = "ltx-2") - is None - ) + assert quantize_text_encoders(pipe, _target(), mode = "fp8_dynamic", family = "ltx-2") is None assert calls == []