lint: pin ruff to the essential ruleset (fixes pre-existing red CI)

docs/workflows/autoencoder.md

@@ -8,7 +8,9 @@ Stable Audio 3 uses a 44.1k stereo audio autoencoder known as SAME to compress w
 import torchaudio
 from stable_audio_3 import AutoencoderModel
 
-ae = AutoencoderModel.from_pretrained("same-l")  # "same-s" (small), "same-l" (medium/large)
+ae = AutoencoderModel.from_pretrained(
+    "same-l"
+)  # "same-s" (small), "same-l" (medium/large)
 waveform, sr = torchaudio.load("audio.wav")
 latents = ae.encode(waveform, sr)
 # → (1, latent_dim, latent_time)

docs/workflows/inference.md

@@ -8,7 +8,10 @@ An overview of the different inference modes. The python interface is shown, but
 
 ```python
 from stable_audio_3 import StableAudioModel
-model = StableAudioModel.from_pretrained("medium", device="cuda")  # device is optional, defaults to cuda → mps → cpu
+
+model = StableAudioModel.from_pretrained(
+    "medium", device="cuda"
+)  # device is optional, defaults to cuda → mps → cpu
 ```
 
 The first argument selects the model to load. Available models:
@@ -34,10 +37,10 @@ audio = model.generate(
     prompt="An anthemic Pop Rock instrumental that fills your head with nostalgic thoughtfulness",
     negative_prompt="poor quality",
     duration=30,
-    steps=8, # default
-    cfg_scale=1, # default
-    seed=-1, # default
-    batch_size=1 # default
+    steps=8,  # default
+    cfg_scale=1,  # default
+    seed=-1,  # default
+    batch_size=1,  # default
 )
 ```
 
@@ -117,7 +120,7 @@ import torchaudio
 from stable_audio_3 import StableAudioModel
 
 model = StableAudioModel.from_pretrained("medium")
-inpaint_audio = torchaudio.load("/path/to/some/audio.wav") # Assume this is 10s long
+inpaint_audio = torchaudio.load("/path/to/some/audio.wav")  # Assume this is 10s long
 audio = model.generate(
     inpaint_audio=inpaint_audio,
     inpaint_mask_start_seconds=10.0,
@@ -211,9 +214,9 @@ model.load_lora(["style_a.safetensors", "style_b.safetensors"])
 Control how strongly the LoRA influences the output at runtime:
 
 ```python
-model.set_lora_strength(0.5)              # Half-strength on all LoRAs
-model.set_lora_strength(1.5)              # Amplify the effect
-model.set_lora_strength(0.0)              # Disable without unloading
+model.set_lora_strength(0.5)  # Half-strength on all LoRAs
+model.set_lora_strength(1.5)  # Amplify the effect
+model.set_lora_strength(0.0)  # Disable without unloading
 
 # With multiple LoRAs, target by index:
 model.set_lora_strength(1.0, lora_index=0)

docs/workflows/lora.md

@@ -309,9 +309,9 @@ The `set_lora_strength()` function adjusts the LoRA contribution at runtime with
 ```python
 from stable_audio_3.models.lora import set_lora_strength
 
-set_lora_strength(model, 0.5)   # Half-strength on all LoRAs
-set_lora_strength(model, 0.0)   # Effectively disable all LoRAs
-set_lora_strength(model, 2.0)   # Double-strength on all LoRAs
+set_lora_strength(model, 0.5)  # Half-strength on all LoRAs
+set_lora_strength(model, 0.0)  # Effectively disable all LoRAs
+set_lora_strength(model, 2.0)  # Double-strength on all LoRAs
 
 # With multiple LoRAs, target a specific one by index:
 set_lora_strength(model, 1.0, lora_index=0)  # Full strength on first LoRA
@@ -328,16 +328,8 @@ You can merge multiple LoRA checkpoints with different weights into a single bas
 from stable_audio_3.models.lora.utils import merge_loras_into_base_model
 
 lora_configurations = [
-    {
-        'name': 'style_a',
-        'state_dict': lora_sd_a,
-        'application_weight': 0.7
-    },
-    {
-        'name': 'style_b',
-        'state_dict': lora_sd_b,
-        'application_weight': 0.3
-    }
+    {"name": "style_a", "state_dict": lora_sd_a, "application_weight": 0.7},
+    {"name": "style_b", "state_dict": lora_sd_b, "application_weight": 0.3},
 ]
 
 merge_loras_into_base_model(model, lora_configurations)
@@ -352,8 +344,8 @@ For models where the input embedding and output projection share weights, LoRA s
 ```python
 from stable_audio_3.models.lora.utils import tie_weights, untie_weights
 
-tie_weights(linear_layer, embedding_layer)   # Share LoRA params
-untie_weights(linear_layer, embedding_layer) # Create independent copies
+tie_weights(linear_layer, embedding_layer)  # Share LoRA params
+untie_weights(linear_layer, embedding_layer)  # Create independent copies
 ```
 
 This is only supported for standard LoRA (not DoRA, BoRA, or -XS variants).

pyproject.toml

@@ -68,3 +68,11 @@ exclude = [
     "stable_audio_3/training",
     "optimized",
 ]
+
+[tool.ruff.lint]
+# Pin the ruleset to the classic essentials. Newer ruff's *default* select is broad
+# (I/UP/B/SIM/PLR/RUF — import-sorting, type-annotation modernization, etc.), which is
+# why CI (ruff >=0.15.9) flagged ~34 style issues across the codebase. This repo only
+# wants real-bug checks, so select them explicitly for version-stable results:
+# E4/E7/E9 (pycodestyle errors) + F (pyflakes: undefined names, unused imports/vars).
+select = ["E4", "E7", "E9", "F"]

tests/test_mlx_latent_dataset.py

@@ -15,14 +15,14 @@
     iterate_batches,
 )
 
-D = 8          # latent channels
-CROP = 16      # latent_crop_length used throughout
+D = 8  # latent channels
+CROP = 16  # latent_crop_length used throughout
 
-LONG_T = 64        # stored length of the long item (> CROP)
-LONG_VALID = 50    # 1s in its padding mask (trailing zeros after)
+LONG_T = 64  # stored length of the long item (> CROP)
+LONG_VALID = 50  # 1s in its padding mask (trailing zeros after)
 LONG_SECONDS = 120.5
 
-SHORT_T = 6        # stored length of the short item (< CROP)
+SHORT_T = 6  # stored length of the short item (< CROP)
 SHORT_SECONDS = 2.786
 
 SILENCE_VALUE = 7.5
@@ -53,14 +53,22 @@ def _make_root(tmp_path, with_silence):
     root = tmp_path / ("latents" if with_silence else "latents_nosilence")
     root.mkdir()
     _write_item(
-        root, "long", LONG_T, LONG_SECONDS,
-        {"title": "Neon Skyline", "artist": "The Testers",
-         "genre": "Synthwave", "bpm": "120"},
+        root,
+        "long",
+        LONG_T,
+        LONG_SECONDS,
+        {
+            "title": "Neon Skyline",
+            "artist": "The Testers",
+            "genre": "Synthwave",
+            "bpm": "120",
+        },
         valid=LONG_VALID,
     )
     # "prompt" tag = what pre_encode extracts from a .txt sidecar caption
-    _write_item(root, "short", SHORT_T, SHORT_SECONDS,
-                {"prompt": "a lofi hip hop beat"})
+    _write_item(
+        root, "short", SHORT_T, SHORT_SECONDS, {"prompt": "a lofi hip hop beat"}
+    )
     if with_silence:
         # Reference stores silence as [1, C, N] (it squeezes axis 0 on load)
         silence = np.full((1, D, 4), SILENCE_VALUE, dtype=np.float32)
@@ -294,8 +302,12 @@ def test_legacy_prompt_without_prompt_config():
 
 def test_legacy_prompt_via_dataset(data_root):
     ds = PreEncodedLatentDataset(data_root, CROP)  # prompt_config=None → legacy
-    expected = {"Artist: The Testers", "Title: Neon Skyline",
-                "BPM: 120", "Genre: Synthwave"}
+    expected = {
+        "Artist: The Testers",
+        "Title: Neon Skyline",
+        "BPM: 120",
+        "Genre: Synthwave",
+    }
     for _ in range(20):
         prompt = _get_by_relpath(ds, "long.npy")["prompt"]
         assert prompt
@@ -308,7 +320,8 @@ def test_txt_derived_prompt_tag(data_root):
     assert _get_by_relpath(ds, "short.npy")["prompt"] == "Prompt: a lofi hip hop beat"
 
     ds = PreEncodedLatentDataset(
-        data_root, CROP,
+        data_root,
+        CROP,
         prompt_config={"shuffle": False, "hide_tag_names": True},
     )
     assert _get_by_relpath(ds, "short.npy")["prompt"] == "a lofi hip hop beat"
@@ -322,6 +335,11 @@ def test_path_prompt_and_space_joined_trigger():
     assert build_prompt(meta, pc, rng) == "artistX/track01"
 
     # non-tag method → trigger joined with a space, not ", "
-    pc = {"use_tags": False, "use_paths": True,
-          "path_opts": {"hideExt": True}, "trigger": "zkq", "trigger_pct": 100}
+    pc = {
+        "use_tags": False,
+        "use_paths": True,
+        "path_opts": {"hideExt": True},
+        "trigger": "zkq",
+        "trigger_pct": 100,
+    }
     assert build_prompt(meta, pc, rng) == "zkq artistX/track01"

tests/test_mlx_lora.py

@@ -534,9 +534,7 @@ def test_trainable_seconds_embedder_matches_pipeline_conditioner():
         assert actual.shape == (len(seconds), 1, 768)
         assert bool(mx.all(actual == expected))
     # mx.array input path (what the trainer passes) is bit-identical too.
-    assert bool(
-        mx.all(trainable(mx.array([12.5, 380.0])) == reference([12.5, 380.0]))
-    )
+    assert bool(mx.all(trainable(mx.array([12.5, 380.0])) == reference([12.5, 380.0])))
 
     report, _ = inject_from_lora_config(
         trainable,
@@ -546,9 +544,7 @@ def test_trainable_seconds_embedder_matches_pipeline_conditioner():
     layer = next(iter(iter_trainable_lora_layers(trainable)))
 
     assert report.layer_names == ("embedder.embedding.1",)
-    assert layer.checkpoint_name == (
-        "conditioners.seconds_total.embedder.embedding.1"
-    )
+    assert layer.checkpoint_name == ("conditioners.seconds_total.embedder.embedding.1")
     # Base Linear is frozen — only the adapter trains.
     assert sorted(
         name for name, _ in tree_flatten(trainable.trainable_parameters())
@@ -711,9 +707,7 @@ def loss_fn(model, values):
 
     new_flat = dict(tree_flatten(grads_new))
     old_flat = dict(tree_flatten(grads_old))
-    expected_params = (
-        {"M_xs"} if adapter_type.endswith("-xs") else {"lora_A", "lora_B"}
-    )
+    expected_params = {"M_xs"} if adapter_type.endswith("-xs") else {"lora_A", "lora_B"}
     if "dora" in adapter_type:
         expected_params = expected_params | {"magnitude"}
     elif "bora" in adapter_type:
@@ -722,9 +716,7 @@ def loss_fn(model, values):
     assert set(new_flat) == set(old_flat)
 
     tol = 1e-3 if dtype == mx.float32 else 3e-2
-    np.testing.assert_allclose(
-        float(loss_new), float(loss_old), rtol=tol, atol=tol
-    )
+    np.testing.assert_allclose(float(loss_new), float(loss_old), rtol=tol, atol=tol)
     for name in sorted(new_flat):
         np.testing.assert_allclose(
             np.asarray(new_flat[name], dtype=np.float32),
@@ -941,9 +933,7 @@ def test_checkpoint_key_naming_matches_real_underfit_checkpoint(tmp_path: Path):
         produced_keys.update(
             f"{root}.{param}" for param in ("lora_A", "lora_B", "magnitude")
         )
-    reference_dit_keys = {
-        key for key in reference_state if key.startswith("model.")
-    }
+    reference_dit_keys = {key for key in reference_state if key.startswith("model.")}
 
     assert produced_keys == reference_dit_keys
     assert report.layer_count == len(reference_dit_keys) // 3
@@ -953,9 +943,7 @@ def test_checkpoint_key_naming_matches_real_underfit_checkpoint(tmp_path: Path):
 
     # The saver reproduces a real medium DiT layer's keys AND shapes exactly.
     single = dit_mlx_medium.DiT(T_lat=8)
-    inject_from_lora_config(
-        single, dict(config, include=["layers.0.self_attn.to_qkv"])
-    )
+    inject_from_lora_config(single, dict(config, include=["layers.0.self_attn.to_qkv"]))
     saved_state, _ = load_lora_checkpoint(
         save_lora_checkpoint(single, tmp_path / "single-layer.safetensors")
     )

tests/test_mlx_pre_encode.py

@@ -10,13 +10,13 @@
 pytest.importorskip("mlx.core")
 soundfile = pytest.importorskip("soundfile")
 
-import mlx.core as mx
+import mlx.core as mx  # noqa: E402
 
 SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "optimized" / "mlx" / "scripts"
 if str(SCRIPTS_DIR) not in sys.path:
     sys.path.insert(0, str(SCRIPTS_DIR))
 
-import pre_encode_mlx as pe
+import pre_encode_mlx as pe  # noqa: E402
 
 SR = 44100
 
@@ -133,9 +133,7 @@ def test_long45_chunked_mono_and_txt_prompt(encoded, dataset):
 
     # >30 s took the chunked path; stub is stride-local → must match unchunked
     audio = pe.load_audio(dataset / "long45.wav")
-    ref = pe.encode_audio(
-        StubEncoder(), audio[None, ...], pad_modulo=32, chunked=False
-    )
+    ref = pe.encode_audio(StubEncoder(), audio[None, ...], pad_modulo=32, chunked=False)
     np.testing.assert_allclose(latents, np.asarray(ref.latents)[0], atol=1e-6)
     assert np.any(latents != 0.0)
 
@@ -192,9 +190,7 @@ def test_encode_file_seconds_rounding_and_mask():
     rng = np.random.default_rng(7)
     audio = rng.standard_normal((2, 130000)).astype(np.float32)
 
-    latents, mask, seconds_total = pe.encode_file(
-        StubEncoder(), audio, pad_modulo=32
-    )
+    latents, mask, seconds_total = pe.encode_file(StubEncoder(), audio, pad_modulo=32)
     assert seconds_total == round(130000 / SR, 3) == 2.948
     # 130000 → padded to 16*8192 = 131072 → 32 latents; ceil(130000/4096) = 32 valid
     assert latents.shape == (8, 32)
@@ -212,7 +208,9 @@ def test_encode_file_seconds_rounding_and_mask():
 
 def test_encode_file_rejects_non_stereo():
     with pytest.raises(ValueError, match=r"\(2, T\)"):
-        pe.encode_file(StubEncoder(), np.zeros((1, 8192), dtype=np.float32), pad_modulo=32)
+        pe.encode_file(
+            StubEncoder(), np.zeros((1, 8192), dtype=np.float32), pad_modulo=32
+        )
 
 
 # ---------------------------------------------------------------------------
@@ -251,9 +249,7 @@ def test_load_audio_resamples_to_44100(tmp_path):
     assert audio.shape == (2, SR)
     # Still a ~110 Hz sine after resampling: correlate with the ideal signal
     ideal = 0.3 * np.sin(2 * np.pi * 110 * np.arange(SR) / SR)
-    corr = np.dot(audio[0], ideal) / (
-        np.linalg.norm(audio[0]) * np.linalg.norm(ideal)
-    )
+    corr = np.dot(audio[0], ideal) / (np.linalg.norm(audio[0]) * np.linalg.norm(ideal))
     assert corr > 0.99