diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 90baaf328c..110e9c34f7 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -687,14 +687,10 @@ def install_hunyuan_attention_trim( if getattr(dit, "_unsloth_trim_hook", None) is None: pre_handle = None try: - pre_handle = dit.register_forward_pre_hook( - _hunyuan_trim_pre_hook, with_kwargs = True - ) + pre_handle = dit.register_forward_pre_hook(_hunyuan_trim_pre_hook, with_kwargs = True) # always_call: clear the flag even when the forward raises, so an exception can # never leave the null-mask authorisation latched for a later direct forward. - post_handle = dit.register_forward_hook( - _hunyuan_trim_post_hook, always_call = True - ) + post_handle = dit.register_forward_hook(_hunyuan_trim_post_hook, always_call = True) dit._unsloth_trim_hook = (pre_handle, post_handle) except Exception as exc: # noqa: BLE001 — optimisation only if pre_handle is not None: diff --git a/studio/backend/core/inference/diffusion_cfg_parallel.py b/studio/backend/core/inference/diffusion_cfg_parallel.py index 79729492d5..8c073851c8 100644 --- a/studio/backend/core/inference/diffusion_cfg_parallel.py +++ b/studio/backend/core/inference/diffusion_cfg_parallel.py @@ -558,9 +558,7 @@ def maybe_enable_cfg_parallel( need = weight_bytes + _REPLICA_HEADROOM_BYTES # Filter by the replica's memory need FIRST so a viable heterogeneous GPU is preferred # over an identical GPU too small to hold the replica. - secondary, free, device_match = _pick_secondary_device( - primary_index, min_free_bytes = need - ) + secondary, free, device_match = _pick_secondary_device(primary_index, min_free_bytes = need) if secondary is None: return None, "no queryable secondary CUDA device" if not device_match: @@ -645,9 +643,7 @@ def maybe_enable_cfg_parallel( # the whole parallel run. A cache may engage/toggle on this DiT, so fullgraph stays # off like the loader. max_speed = str(speed_mode) == SPEED_MAX - _compile_repeated_blocks( - view, logger, max_autotune = max_speed, cache_active = True - ) + _compile_repeated_blocks(view, logger, max_autotune = max_speed, cache_active = True) if max_speed: # Fuse the REPLICA's QKV projections directly: _fuse_qkv(view) would resolve the # pipe-level fuse_qkv_projections through _ReplicaView delegation and re-fuse the diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index 58158ee9fc..b1ac0e429b 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -340,6 +340,7 @@ def quantize_text_encoders( # upcast hooks + fp8 storage in place, leaving no torchao params). Detect a leftover # layerwise hook directly and fail closed there too; a clean failure stays dense. from .diffusion_transformer_quant import raise_if_partially_quantized + if mode == TE_QUANT_FP8 and _has_layerwise_casting(encoder): raise RuntimeError( f"text_encoder_quant fp8:{attr} failed after partially installing layerwise " diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index b521eabc17..98e88aec73 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -2189,9 +2189,7 @@ class VideoBackend: # Transactional across MoE experts (like the load / AUTO paths): refuse to # stack a fresh cache over one whose removal failed, and roll back a mixed # resize so status never reports MagCache over an asymmetric pair. - def _resize_explicit_magcache( - view: Any, expert_name: str - ) -> Optional[str]: + def _resize_explicit_magcache(view: Any, expert_name: str) -> Optional[str]: transformer = getattr(view, "transformer", None) marker = getattr(transformer, "_unsloth_step_cache", None) # endswith, not substring: "#s5" would match inside "#s50". diff --git a/studio/backend/tests/test_diffusion_cfg_parallel.py b/studio/backend/tests/test_diffusion_cfg_parallel.py index 6dd623387b..231f859060 100644 --- a/studio/backend/tests/test_diffusion_cfg_parallel.py +++ b/studio/backend/tests/test_diffusion_cfg_parallel.py @@ -659,13 +659,19 @@ def test_replica_mirrors_max_tier_compile_and_fusion(monkeypatch): speed = _engage_stubs(monkeypatch) compile_kwargs: list = [] monkeypatch.setattr( - speed, "_compile_repeated_blocks", + speed, + "_compile_repeated_blocks", lambda view, logger, **kw: compile_kwargs.append(kw) or True, ) pipe = _CtxPipe(_LoadableFuseDiT()) proxy, reason = _gate( - monkeypatch, pipe, _fam(), requested = "on", compiled = True, - speed_mode = "max", attention_backend = None, + monkeypatch, + pipe, + _fam(), + requested = "on", + compiled = True, + speed_mode = "max", + attention_backend = None, ) assert proxy is not None, reason try: @@ -682,13 +688,19 @@ def test_replica_default_tier_no_max_autotune_no_fuse(monkeypatch): speed = _engage_stubs(monkeypatch) compile_kwargs: list = [] monkeypatch.setattr( - speed, "_compile_repeated_blocks", + speed, + "_compile_repeated_blocks", lambda view, logger, **kw: compile_kwargs.append(kw) or True, ) pipe = _CtxPipe(_LoadableFuseDiT()) proxy, reason = _gate( - monkeypatch, pipe, _fam(), requested = "on", compiled = True, - speed_mode = "default", attention_backend = None, + monkeypatch, + pipe, + _fam(), + requested = "on", + compiled = True, + speed_mode = "default", + attention_backend = None, ) assert proxy is not None, reason try: @@ -709,12 +721,17 @@ def test_replica_pins_native_when_backend_unsupported_on_secondary(monkeypatch): monkeypatch.setattr(attn, "attention_backend_supported_on_device", lambda backend, idx: False) applied: list = [] monkeypatch.setattr( - attn, "apply_attention_backend", + attn, + "apply_attention_backend", lambda view, backend, logger = None: applied.append(backend) or backend, ) pipe = _CtxPipe(_LoadableFuseDiT()) proxy, reason = _gate( - monkeypatch, pipe, _fam(), requested = "on", attention_backend = "_flash_3_hub", + monkeypatch, + pipe, + _fam(), + requested = "on", + attention_backend = "_flash_3_hub", ) assert proxy is not None, reason try: diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 7ba097fc10..0a1e7083a1 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -2148,12 +2148,14 @@ def test_step_cache_all_or_none_raises_when_rollback_fails(monkeypatch): pipe, fam, _t1, _t2 = _moe_pipe_and_fam() monkeypatch.setattr( - video, "_disengage_step_cache", + video, + "_disengage_step_cache", lambda transformer, *, reason, logger = None: False, # rollback fails ) with pytest.raises(RuntimeError, match = "rollback failed"): video._step_cache_all_or_none( - pipe, fam, + pipe, + fam, lambda view, expert_name: "fbcache" if expert_name == "transformer" else None, logger = None, ) @@ -2174,7 +2176,8 @@ def test_explicit_magcache_hard_errors_when_disable_fails(fake_runtime, monkeypa reapplied: list = [] monkeypatch.setattr(video, "_disengage_step_cache", lambda *a, **k: False) monkeypatch.setattr( - video, "apply_step_cache", + video, + "apply_step_cache", lambda *a, **k: reapplied.append(k.get("steps")) or "magcache", ) with pytest.raises(RuntimeError, match = "reload the video model"):