ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-07-10 14:30:38 +00:00
commit d879a90bfa
8 changed files with 118 additions and 65 deletions
Download patch file Download diff file Expand all files Collapse all files

3
scripts/video_speedmem_bench.py
View file

@ -574,8 +574,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),
logger = logger,
)
except Exception:

1
studio/backend/core/inference/diffusion_cache.py
View file

@ -119,6 +119,7 @@ def auto_cache_quality(family: Optional[str]) -> str:
"""The cache quality preset an UNSET request resolves to for ``family``."""
return _FAMILY_AUTO_CACHE_QUALITY.get(str(family or "").strip().lower(), CQ_BALANCED)
# The auto policy's step-count bar: FBCache's win scales with step count (each skipped
# step is a larger quality hit on a short trajectory), so auto engages it only at 20+
# steps -- full "dev"-style schedules (28+) qualify, distilled turbo models (4-9) never do.

59
studio/backend/core/inference/diffusion_cfg_parallel.py
View file

@ -197,7 +197,11 @@ class CFGParallelProxy:
orig_forward = self._orig_guider_forward
def _device_homogenising_forward(pred_cond, pred_uncond = None, **kw):
def _device_homogenising_forward(
pred_cond,
pred_uncond = None,
**kw,
):
return orig_forward(_resolve(pred_cond), _resolve(pred_uncond), **kw)
guider.forward = _device_homogenising_forward
@ -321,7 +325,6 @@ class CFGParallelProxy:
def _worker_loop(self) -> None:
import torch
while True:
job = self._jobs.get()
if job is None: # shutdown sentinel
@ -372,16 +375,31 @@ def _install_threadsafe_cudnn_attention(logger: Any = None) -> bool:
orig = ad._native_cudnn_attention
def _threadsafe_cudnn_attention(
query, key, value, attn_mask = None, dropout_p = 0.0, is_causal = False,
scale = None, enable_gqa = False, return_lse = False, _parallel_config = None,
query,
key,
value,
attn_mask = None,
dropout_p = 0.0,
is_causal = False,
scale = None,
enable_gqa = False,
return_lse = False,
_parallel_config = None,
):
# Fall back to the stock (context-managed) path for the shapes/options the
# direct op does not cover; the video DiT hot path never takes them.
if _parallel_config is not None or return_lse or enable_gqa or dropout_p:
return orig(
query, key, value, attn_mask = attn_mask, dropout_p = dropout_p,
is_causal = is_causal, scale = scale, enable_gqa = enable_gqa,
return_lse = return_lse, _parallel_config = _parallel_config,
query,
key,
value,
attn_mask = attn_mask,
dropout_p = dropout_p,
is_causal = is_causal,
scale = scale,
enable_gqa = enable_gqa,
return_lse = return_lse,
_parallel_config = _parallel_config,
)
q, k, v = (x.permute(0, 2, 1, 3).contiguous() for x in (query, key, value))
out = torch.ops.aten._scaled_dot_product_cudnn_attention(
@ -513,12 +531,16 @@ def maybe_enable_cfg_parallel(
# ── build the replica and mirror the primary's levers ──
try:
replica = type(primary).from_pretrained(
transformer_source,
subfolder = "transformer",
torch_dtype = dtype,
token = hf_token or None,
).to(f"cuda:{secondary}")
replica = (
type(primary)
.from_pretrained(
transformer_source,
subfolder = "transformer",
torch_dtype = dtype,
token = hf_token or None,
)
.to(f"cuda:{secondary}")
)
replica.eval()
except Exception as exc: # noqa: BLE001 -- download/VRAM race: stay single-device
_warn(logger, "cfg-parallel replica load", exc)
@ -571,7 +593,11 @@ def maybe_enable_cfg_parallel(
return proxy, f"engaged: DiT replica on cuda:{secondary}"
def teardown_cfg_parallel(pipe: Any, proxy: Any, logger: Any = None) -> None:
def teardown_cfg_parallel(
pipe: Any,
proxy: Any,
logger: Any = None,
) -> None:
"""Restore the pipe to its single-device shape and free the replica's VRAM.
Safe to call with a half-built or foreign object; never raises."""
try:
@ -601,10 +627,11 @@ def teardown_cfg_parallel(pipe: Any, proxy: Any, logger: Any = None) -> None:
def _invalidate_registry(module: Any) -> None:
from .diffusion_cache import _invalidate_child_registry_cache
_invalidate_child_registry_cache(module)
def _warn(logger: Any, what: str, exc: Exception) -> None:
if logger is not None:
logger.warning("diffusion.cfg_parallel: %s unavailable (%s); running single-device", what, exc)
logger.warning(
"diffusion.cfg_parallel: %s unavailable (%s); running single-device", what, exc
)

4
studio/backend/core/inference/diffusion_precision.py
View file

@ -135,9 +135,7 @@ _TE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {}
# measured out-of-bar trajectory drift for zero speed win on the video families below.
# Unlike the deny table this only steers the AUTO default; an explicit scheme request
# (text_encoder_quant="fp8_dynamic") is still honored verbatim.
_TE_AUTO_DENSE_FAMILIES: frozenset[str] = frozenset(
{"hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"}
)
_TE_AUTO_DENSE_FAMILIES: frozenset[str] = frozenset({"hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"})
# Map a TE torchao scheme to the transformer smoke-probe scheme (same torchao GEMM), so ``auto``
# degrades gracefully when a build lacks a kernel. Layerwise fp8 has no torchao GEMM to probe.

2
studio/backend/core/inference/diffusion_speed.py
View file

@ -116,7 +116,6 @@ def _inductor_config() -> Any:
reports None instead of picking a stale real module out of ``sys.modules``."""
try:
import torch
return getattr(getattr(torch, "_inductor", None), "config", None)
except Exception: # noqa: BLE001 — no inductor -> nothing to snapshot/set
return None
@ -392,7 +391,6 @@ def _compile_repeated_blocks(
# this module, but keep the dependency one-directional at import time.
try:
from .diffusion_cache import _compile_hooked_block_inners
_compile_hooked_block_inners(transformer, logger)
except Exception as exc: # noqa: BLE001 — optimisation only
_warn(logger, "cache-hook inner compile", exc)

29
studio/backend/tests/test_diffusion_cache.py
View file

@ -666,7 +666,10 @@ def test_magcache_quality_preset_engages_conservative_params(monkeypatch):
_stub_diffusers_with_magcache(monkeypatch)
t = _MixinTransformer()
engaged = apply_step_cache(
_pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p", steps = 50,
_pipe(t),
mode = "magcache",
family = "hunyuanvideo-1.5-720p",
steps = 50,
quality = "quality",
)
assert engaged == TC_MAGCACHE
@ -682,8 +685,12 @@ def test_magcache_explicit_threshold_beats_the_preset(monkeypatch):
_stub_diffusers_with_magcache(monkeypatch)
t = _MixinTransformer()
apply_step_cache(
_pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p", steps = 50,
quality = "fast", threshold = 0.05,
_pipe(t),
mode = "magcache",
family = "hunyuanvideo-1.5-720p",
steps = 50,
quality = "fast",
threshold = 0.05,
)
assert t.enabled_with.threshold == 0.05
assert t.enabled_with.max_skip_steps == _MAGCACHE_QUALITY_PRESETS[CQ_FAST][1]
@ -710,7 +717,10 @@ def test_toggle_threads_quality_through(monkeypatch):
_stub_diffusers_with_magcache(monkeypatch)
t = _ToggleTransformer()
maybe_toggle_step_cache(
_pipe(t), steps = 30, mode = TC_MAGCACHE, family = "hunyuanvideo-1.5-720p",
_pipe(t),
steps = 30,
mode = TC_MAGCACHE,
family = "hunyuanvideo-1.5-720p",
quality = "quality",
)
assert t.enabled_with.threshold == _MAGCACHE_QUALITY_PRESETS[CQ_QUALITY][0]
@ -741,7 +751,12 @@ class _BoundInner:
return "eager"
def _hooked_block(*, compiled = True, hook_name = "mag_cache_block_hook", bound = True):
def _hooked_block(
*,
compiled = True,
hook_name = "mag_cache_block_hook",
bound = True,
):
inner = _BoundInner()
orig = inner.forward if bound else functools.partial(_BoundInner.forward, inner)
hook = types.SimpleNamespace(fn_ref = types.SimpleNamespace(original_forward = orig))
@ -869,9 +884,7 @@ def test_apply_step_cache_arms_compiled_blocks_on_toggle(monkeypatch):
return [block]
t = _T()
engaged = apply_step_cache(
_pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p", steps = 50
)
engaged = apply_step_cache(_pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p", steps = 50)
assert engaged == TC_MAGCACHE
assert hook.fn_ref.original_forward is not orig
assert hook._unsloth_orig_inner is orig

78
studio/backend/tests/test_diffusion_cfg_parallel.py
View file

@ -47,7 +47,11 @@ def test_normalize_rejects_unknown():
# ── fakes ─────────────────────────────────────────────────────────────────────────
class _FakeDevice:
def __init__(self, type_ = "cuda", index = 0):
def __init__(
self,
type_ = "cuda",
index = 0,
):
self.type = type_
self.index = index
@ -55,7 +59,12 @@ class _FakeDevice:
class _FakeTensor:
"""Just enough tensor for the proxy's _move / guider resolve paths."""
def __init__(self, device, tag = "t", nbytes = 8):
def __init__(
self,
device,
tag = "t",
nbytes = 8,
):
self.device = device
self.tag = tag
self._nbytes = nbytes
@ -66,12 +75,20 @@ class _FakeTensor:
def element_size(self):
return 1
def to(self, device, non_blocking = False):
def to(
self,
device,
non_blocking = False,
):
return _FakeTensor(device, tag = self.tag, nbytes = self._nbytes)
class _FakeDiT:
def __init__(self, device_index = 0, fail_enable = False):
def __init__(
self,
device_index = 0,
fail_enable = False,
):
self._device = _FakeDevice(index = device_index)
self.fail_enable = fail_enable
self.enabled_with = None
@ -82,9 +99,9 @@ class _FakeDiT:
self._mods = [self, types.SimpleNamespace(name = f"block{device_index}")]
def parameters(self):
return iter([types.SimpleNamespace(
numel = lambda: 100, element_size = lambda: 2, device = self._device
)])
return iter(
[types.SimpleNamespace(numel = lambda: 100, element_size = lambda: 2, device = self._device)]
)
def modules(self):
return list(self._mods)
@ -110,7 +127,12 @@ class _FakeDiT:
return (_FakeTensor(self._device, tag = "pred"),)
def _stub_torch(monkeypatch, *, device_count = 2, free = None):
def _stub_torch(
monkeypatch,
*,
device_count = 2,
free = None,
):
torch = types.ModuleType("torch")
torch.Tensor = _FakeTensor
free = free if free is not None else {}
@ -129,14 +151,18 @@ def _stub_torch(monkeypatch, *, device_count = 2, free = None):
return torch
def _make_proxy(monkeypatch, *, compiled = False, explicit_on = False, fail_enable = False):
def _make_proxy(
monkeypatch,
*,
compiled = False,
explicit_on = False,
fail_enable = False,
):
_stub_torch(monkeypatch)
primary = _FakeDiT(device_index = 0)
replica = _FakeDiT(device_index = 1, fail_enable = fail_enable)
guider = types.SimpleNamespace(forward = lambda *a, **k: ("combined", a, k), num_conditions = 2)
proxy = CFGParallelProxy(
primary, replica, guider, compiled = compiled, explicit_on = explicit_on
)
proxy = CFGParallelProxy(primary, replica, guider, compiled = compiled, explicit_on = explicit_on)
return proxy, primary, replica, guider
@ -294,9 +320,7 @@ def test_replica_enable_failure_reraises_and_breaks(monkeypatch):
with pytest.raises(RuntimeError):
proxy.enable_cache({})
assert primary.enabled_with == {} # primary was hooked before the replica failed
plan = proxy.plan_generation(
cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33
)
plan = proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33)
assert plan["enabled"] is False
proxy.shutdown()
@ -358,14 +382,10 @@ def test_thread_dispatch_resolves_through_guider(monkeypatch):
# ── per-generation dispatch policy ──────────────────────────────────────────────────
def test_plan_parallel_on_eager_settles_to_thread(monkeypatch):
proxy, _, _, _ = _make_proxy(monkeypatch, compiled = False)
plan = proxy.plan_generation(
cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33
)
plan = proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33)
assert plan["enabled"] is True and plan["dispatch"] == "inline" # first run: compile-safe
proxy.note_generation_done()
plan = proxy.plan_generation(
cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33
)
plan = proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33)
assert plan["dispatch"] == "thread" # settled key: full overlap
proxy.shutdown()
@ -391,9 +411,7 @@ def test_plan_sequential_for_compiled_stack_even_with_cache(monkeypatch):
def test_plan_parallel_for_eager_stack(monkeypatch):
proxy, _, _, _ = _make_proxy(monkeypatch, compiled = False)
plan = proxy.plan_generation(
cache_engaged = False, steps = 10, width = 1280, height = 720, frames = 33
)
plan = proxy.plan_generation(cache_engaged = False, steps = 10, width = 1280, height = 720, frames = 33)
assert plan["enabled"] is True and plan["lossless"] is True
proxy.shutdown()
@ -402,9 +420,7 @@ def test_plan_requires_cfg_conditions(monkeypatch):
# guidance ~1 collapses the guider to one condition: nothing to overlap.
proxy, _, _, guider = _make_proxy(monkeypatch)
guider.num_conditions = 1
plan = proxy.plan_generation(
cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33
)
plan = proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33)
assert plan["enabled"] is False
proxy.shutdown()
@ -413,9 +429,7 @@ def test_shape_change_forces_inline_once(monkeypatch):
proxy, _, _, _ = _make_proxy(monkeypatch)
proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33)
proxy.note_generation_done()
plan = proxy.plan_generation(
cache_engaged = True, steps = 30, width = 960, height = 544, frames = 33
)
plan = proxy.plan_generation(cache_engaged = True, steps = 30, width = 960, height = 544, frames = 33)
assert plan["dispatch"] == "inline" # new shape may recompile: serialize
proxy.shutdown()
@ -424,9 +438,7 @@ def test_cancelled_generation_stays_inline(monkeypatch):
proxy, _, _, _ = _make_proxy(monkeypatch)
proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33)
# No note_generation_done (cancel/failure): the same key must stay inline.
plan = proxy.plan_generation(
cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33
)
plan = proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33)
assert plan["dispatch"] == "inline"
proxy.shutdown()

7
studio/backend/tests/test_diffusion_speed.py
View file

@ -536,7 +536,12 @@ def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch):
# ── inductor precision-cast emulation (compile-vs-eager numeric parity) ─────────
def _stub_inductor_config(monkeypatch, torch, *, emulate = False):
def _stub_inductor_config(
monkeypatch,
torch,
*,
emulate = False,
):
"""Attach a fake ``_inductor.config`` to the stubbed torch module (diffusion_speed
resolves it as attributes off the imported torch, never via sys.modules -- so the
real torch._inductor lingering in sys.modules cannot leak into stubbed tests)."""