lint: ruff fixes — __all__ re-exports, rename ambiguous test vars, format

Author: “Cortexelus”

stable_audio_3/utils/__init__.py

@@ -5,3 +5,11 @@
     make_grad_scaler,
     resolve_device,
 )
+
+__all__ = [
+    "autocast_context",
+    "disable_autocast",
+    "empty_device_cache",
+    "make_grad_scaler",
+    "resolve_device",
+]

stable_audio_3/utils/device.py

@@ -12,6 +12,7 @@
   - @autocast("cuda", enabled=False) does NOT disable an active MPS autocast,
     so fp32 islands need the device-aware `disable_autocast` below.
 """
+
 import functools
 
 import torch

tests/test_fast_lora_forward.py

@@ -5,6 +5,7 @@
 forward: same outputs, same gradients, same state_dict layout, and graceful
 fallback everywhere the reformulation doesn't apply.
 """
+
 import copy
 from functools import partial
 
@@ -64,7 +65,9 @@ def randomize_lora(model, seed=1, magnitude_jitter=0.3):
             for name in ("magnitude",):
                 if hasattr(mod, name):
                     p = getattr(mod, name)
-                    p.mul_(1.0 + magnitude_jitter * torch.rand(p.shape, generator=g).to(p))
+                    p.mul_(
+                        1.0 + magnitude_jitter * torch.rand(p.shape, generator=g).to(p)
+                    )
 
 
 def lora_trainables(model):
@@ -147,7 +150,9 @@ def test_gradient_parity_mps_fp16_autocast(adapter):
         y_naive, g_naive = run_fwd_bwd(naive, x, proj)
     out_err = rel_err(y_fast.float(), y_naive.float())
     worst = max(rel_err(g_fast[k].float(), g_naive[k].float()) for k in g_naive)
-    print(f"[autocast fp16 {adapter}] output rel err = {out_err:.3e}, worst grad rel err = {worst:.3e}")
+    print(
+        f"[autocast fp16 {adapter}] output rel err = {out_err:.3e}, worst grad rel err = {worst:.3e}"
+    )
     # fp16 GEMM noise dominates; ~1e-2-class is the expected scale, this bound is a regression guard
     assert out_err <= 3e-2
     assert worst <= 5e-2
@@ -166,7 +171,9 @@ def test_strength_buffer(adapter, device):
     base = make_model(bias=True, device=device)  # same seed -> identical base weights
     y_fast = fast(x)
     y_base = base(x)
-    assert torch.equal(y_fast, y_base), f"strength=0 must be bit-exact to the base ({adapter}/{device})"
+    assert torch.equal(y_fast, y_base), (
+        f"strength=0 must be bit-exact to the base ({adapter}/{device})"
+    )
 
 
 @pytest.mark.parametrize("device", DEVICES)
@@ -179,9 +186,12 @@ def test_stacked_parametrization_falls_back(device):
     for m in (fast, naive):
         torch.manual_seed(7)
         for mod in m.modules():
-            if isinstance(mod, nn.Linear) and parametrize.is_parametrized(mod, "weight"):
+            if isinstance(mod, nn.Linear) and parametrize.is_parametrized(
+                mod, "weight"
+            ):
                 p2 = LoRAParametrization.from_linear(
-                    mod, rank=4, lora_alpha=8, adapter_type="dora-rows", lora_index=1)
+                    mod, rank=4, lora_alpha=8, adapter_type="dora-rows", lora_index=1
+                )
                 with torch.no_grad():
                     p2.lora_B.copy_(torch.randn_like(p2.lora_B) * 0.05)
                 parametrize.register_parametrization(mod, "weight", p2, unsafe=True)
@@ -192,7 +202,7 @@ def test_stacked_parametrization_falls_back(device):
             assert len(mod.parametrizations["weight"]) == 2
 
     x = torch.randn(2, 9, 64, device=device)
-    y_fast = fast(x)     # wrapper must detect the stack and fall back per-forward
+    y_fast = fast(x)  # wrapper must detect the stack and fall back per-forward
     y_naive = naive(x)
     assert torch.equal(y_fast, y_naive), "stacked modules must use the exact naive path"
 
@@ -229,7 +239,9 @@ def test_state_dict_layout_unchanged(adapter, monkeypatch):
     monkeypatch.setenv("SA3_FAST_LORA", "0")
     naive = make_model()
     add_lora(naive, lora_config(adapter))
-    assert not naive[0].__dict__.get("_fast_lora_wrapped", False), "SA3_FAST_LORA=0 must disable wrapping"
+    assert not naive[0].__dict__.get("_fast_lora_wrapped", False), (
+        "SA3_FAST_LORA=0 must disable wrapping"
+    )
 
     # populate the norm-constant cache before snapshotting
     fast(torch.randn(2, 4, 64))
@@ -249,8 +261,13 @@ def test_dropout_shared_between_direction_and_norm(device):
     forward; under a fixed RNG seed both paths must consume the same draws."""
     cfg = {
         nn.Linear: {
-            "weight": partial(LoRAParametrization.from_linear, rank=8, lora_alpha=16,
-                              adapter_type="dora-rows", lora_dropout_p=0.5),
+            "weight": partial(
+                LoRAParametrization.from_linear,
+                rank=8,
+                lora_alpha=16,
+                adapter_type="dora-rows",
+                lora_dropout_p=0.5,
+            ),
         },
     }
     model = make_model(device=device)

tests/test_mps_training_smoke.py

@@ -15,6 +15,7 @@
 Skipped entirely when MPS is unavailable. Runnable standalone:
     python tests/test_mps_training_smoke.py
 """
+
 import os
 import sys
 from functools import partial
@@ -36,19 +37,25 @@
 # underfit loss (real module when importable, faithful mirror otherwise)
 # ---------------------------------------------------------------------------
 
+
 def _underfit_loss_fns():
     try:
         from underfit.training.loss import compute_masked_loss, compute_normalized_mse
+
         return compute_normalized_mse, compute_masked_loss
     except ImportError:
         pass
 
     # Mirror of underfit/training/loss.py (loss_normalization="none",
     # mask_padding_attention=True → signal-only masked MSE).
-    def compute_normalized_mse(pred, target, loss_mask, loss_normalization="none", loss_norm_eps=1e-6):
+    def compute_normalized_mse(
+        pred, target, loss_mask, loss_normalization="none", loss_norm_eps=1e-6
+    ):
         return (pred - target) ** 2
 
-    def compute_masked_loss(loss_full, loss_mask, mask_padding_attention, mask_loss_weight=0.0):
+    def compute_masked_loss(
+        loss_full, loss_mask, mask_padding_attention, mask_loss_weight=0.0
+    ):
         signal = torch.where(loss_mask.unsqueeze(1), loss_full, 0.0)
         signal_sum = signal.sum(dim=(1, 2))
         n_channels = loss_full.shape[1]
@@ -65,6 +72,7 @@ def compute_masked_loss(loss_full, loss_mask, mask_padding_attention, mask_loss_
 # Device helper probes
 # ---------------------------------------------------------------------------
 
+
 def test_resolve_device_prefers_mps_without_cuda():
     from stable_audio_3.utils.device import resolve_device
 
@@ -120,6 +128,7 @@ def test_fp32_islands_hold_under_mps_autocast():
 # End-to-end: tiny DiT + LoRA, 3 training steps on MPS
 # ---------------------------------------------------------------------------
 
+
 def _build_tiny_dit_with_lora():
     from stable_audio_3.models.dit import DiffusionTransformer
     from stable_audio_3.models.lora import LoRAParametrization, add_lora
@@ -206,7 +215,9 @@ def test_lora_training_steps_on_mps():
         scaler.update()
         losses.append(loss.item())
 
-    assert all(l == l and abs(l) != float("inf") for l in losses), f"non-finite loss: {losses}"
+    assert all(x == x and abs(x) != float("inf") for x in losses), (
+        f"non-finite loss: {losses}"
+    )
     assert losses[-1] != losses[0], f"loss did not change over 3 steps: {losses}"
     # On a fixed batch with lr=1e-2 the loss should trend down.
     assert losses[-1] < losses[0] * 1.05, f"loss did not decrease: {losses}"
@@ -228,7 +239,10 @@ def test_full_wrapper_training_step_on_mps():
     (underfit's pre_encoded path never calls pretransform.encode).
     """
     from stable_audio_3.models.conditioners import MultiConditioner, NumberConditioner
-    from stable_audio_3.models.diffusion import ConditionedDiffusionModelWrapper, DiTWrapper
+    from stable_audio_3.models.diffusion import (
+        ConditionedDiffusionModelWrapper,
+        DiTWrapper,
+    )
     from stable_audio_3.models.lora import LoRAParametrization, add_lora
     from stable_audio_3.utils.device import autocast_context, make_grad_scaler
 
@@ -310,7 +324,9 @@ def test_full_wrapper_training_step_on_mps():
         scaler.update()
         losses.append(loss.item())
 
-    assert all(l == l and abs(l) != float("inf") for l in losses), f"non-finite loss: {losses}"
+    assert all(x == x and abs(x) != float("inf") for x in losses), (
+        f"non-finite loss: {losses}"
+    )
     assert len(set(losses)) > 1, f"loss frozen across steps: {losses}"