diff --git a/docs/workflows/autoencoder.md b/docs/workflows/autoencoder.md index 6bf95569..2b2190a2 100644 --- a/docs/workflows/autoencoder.md +++ b/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) diff --git a/docs/workflows/inference.md b/docs/workflows/inference.md index 4988f287..aa2d6d14 100644 --- a/docs/workflows/inference.md +++ b/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) diff --git a/docs/workflows/lora.md b/docs/workflows/lora.md index e992a715..83472830 100644 --- a/docs/workflows/lora.md +++ b/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). diff --git a/pyproject.toml b/pyproject.toml index 6a48191f..9dd1797f 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/test_mlx_latent_dataset.py b/tests/test_mlx_latent_dataset.py index 5830072c..dabd578a 100644 --- a/tests/test_mlx_latent_dataset.py +++ b/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" diff --git a/tests/test_mlx_lora.py b/tests/test_mlx_lora.py index d8e8e247..3be8d60e 100644 --- a/tests/test_mlx_lora.py +++ b/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") ) diff --git a/tests/test_mlx_pre_encode.py b/tests/test_mlx_pre_encode.py index 1e0a07de..5fa42a12 100644 --- a/tests/test_mlx_pre_encode.py +++ b/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