diff --git a/src/diffusers/schedulers/sway_sampling.py b/src/diffusers/schedulers/sway_sampling.py new file mode 100644 index 000000000000..4ff46958ced1 --- /dev/null +++ b/src/diffusers/schedulers/sway_sampling.py @@ -0,0 +1,134 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Inference-time *sway sampling* for flow-matching schedulers. + +Implements the flow-step sampling strategy from F5-TTS (Chen et al., 2024), which +redistributes the uniformly spaced flow steps used at inference time via a +monotone map over the [0, 1] progress coordinate. Concentrating steps toward one +end of the denoising trajectory gives the ODE solver finer resolution where it +matters most, improving the quality/efficiency trade-off **without retraining** -- +the paper notes the strategy "can be easily applied to existing flow matching +based models without retraining." + +The mapping (F5-TTS, Sec. 3.2) is:: + + f_sway(u; s) = u + s * (cos(pi/2 * u) - 1 + u) , u in [0, 1] + +``f_sway(0) = 0`` and ``f_sway(1) = 1`` for every ``s``, and the map is monotone on +``s in [-1, 2 / (pi - 2)]``: + +* ``s < 0`` -> steps are pulled toward the *start* of the trajectory (high noise + level), giving the solver finer resolution on the early integration steps. +* ``s > 0`` -> steps are pulled toward the *end* of the trajectory. +* ``s = 0`` -> recovers the original (uniform) schedule. + +F5-TTS uses ``s = -1`` by default; that is the default here too. + +Reference: "F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow +Matching", https://arxiv.org/abs/2410.06885 +""" + +import math + +import numpy as np +import torch + + +# Monotonicity interval for the sway coefficient, as derived in the F5-TTS paper. +SWAY_COEF_MIN = -1.0 +SWAY_COEF_MAX = 2.0 / (math.pi - 2.0) +# F5-TTS default (paper Sec. 5). +DEFAULT_SWAY_COEF = -1.0 + + +def sway_sampling(u, sway_coef=DEFAULT_SWAY_COEF): + """Map flow-step positions through F5-TTS *sway sampling*. + + Args: + u (`np.ndarray` or `torch.Tensor`): + Flow-step positions in `[0, 1]` -- typically a uniform grid spanning + the inference steps. The endpoints `0` and `1` are fixed points. + sway_coef (`float`, defaults to `-1.0`): + Sway coefficient `s`. Must lie in `[-1, 2 / (pi - 2)]` so the map + stays monotone. Negative values bias steps toward the start of the + trajectory, positive values toward the end, and `0` leaves the grid + unchanged. + + Returns: + An array of the same type and shape as `u` holding the sway-mapped + positions in `[0, 1]`. + + Raises: + ValueError: if `sway_coef` falls outside its monotone interval. + """ + if not (SWAY_COEF_MIN - 1e-9 <= sway_coef <= SWAY_COEF_MAX + 1e-9): + raise ValueError( + f"`sway_coef` must be in [{SWAY_COEF_MIN:.4f}, {SWAY_COEF_MAX:.4f}] to keep sway sampling monotone, " + f"got {sway_coef}." + ) + # f_sway(u; s) = u + s * (cos(pi/2 * u) - 1 + u). Dispatch on the backend so + # the function works for both numpy arrays and torch tensors. + if isinstance(u, torch.Tensor): + half_pi = math.pi / 2.0 + return u + sway_coef * (torch.cos(half_pi * u) - 1.0 + u) + return u + sway_coef * (np.cos(np.pi / 2.0 * u) - 1.0 + u) + + +def apply_sway_sampling(scheduler, sway_coef=DEFAULT_SWAY_COEF): + """Re-space a flow-matching scheduler's schedule via sway sampling. + + Operates on any scheduler that exposes ``sigmas`` (a 1-D tensor of length + ``num_inference_steps + 1`` decreasing from the starting noise level down to + ``0``) and derives ``timesteps`` as ``sigmas * num_train_timesteps`` -- e.g. + [`~diffusers.FlowMatchEulerDiscreteScheduler`]. The already-configured + schedule (shift, Karras, ...) is left intact: sway sampling only re-spaces + *where along the trajectory* each step lands, exactly as F5-TTS applies it at + inference time. + + Call this immediately after ``scheduler.set_timesteps(...)`` and before the + sampling loop. The change is in-place and requires no retraining. + + Args: + scheduler: + A configured flow-matching scheduler with `sigmas`/`timesteps` set. + sway_coef (`float`, defaults to `-1.0`): + Sway coefficient; see [`sway_sampling`]. + """ + sigmas = getattr(scheduler, "sigmas", None) + if sigmas is None: + raise ValueError("Scheduler has no sigma schedule; call `set_timesteps` before applying sway sampling.") + num_nodes = int(sigmas.shape[0]) + # With fewer than three nodes there is no interior point to re-space. + if num_nodes < 3: + return + + device, dtype = sigmas.device, sigmas.dtype + sigmas_np = sigmas.detach().to("cpu", dtype=torch.float32).numpy() + + # Uniform progress coordinate over the schedule nodes (0 = start, 1 = clean). + u = np.linspace(0.0, 1.0, num_nodes, dtype=np.float32) + # Sway-map the progress coordinate and resample the original noise levels at + # the mapped positions; `np.interp` keeps the endpoints fixed by construction. + t = np.asarray(sway_sampling(u, sway_coef), dtype=np.float32) + new_sigmas = np.interp(t, u, sigmas_np).astype(np.float32) + + new_sigmas_t = torch.from_numpy(new_sigmas).to(device=device, dtype=dtype) + num_train_timesteps = scheduler.config.num_train_timesteps + new_timesteps = new_sigmas_t[:-1] * num_train_timesteps + + scheduler.sigmas = new_sigmas_t + scheduler.timesteps = new_timesteps.to(dtype=torch.float32) + scheduler._step_index = None + scheduler._begin_index = None diff --git a/tests/schedulers/test_sway_sampling.py b/tests/schedulers/test_sway_sampling.py new file mode 100644 index 000000000000..8ec8cd27d83e --- /dev/null +++ b/tests/schedulers/test_sway_sampling.py @@ -0,0 +1,102 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import pytest +import torch + +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.schedulers.sway_sampling import ( + DEFAULT_SWAY_COEF, + SWAY_COEF_MAX, + SWAY_COEF_MIN, + apply_sway_sampling, + sway_sampling, +) + + +def test_sway_sampling_formula(): + u = np.linspace(0.0, 1.0, 9, dtype=np.float32) + # s = 0 must be the identity mapping. + np.testing.assert_allclose(sway_sampling(u, 0.0), u, atol=1e-7) + # Endpoints are fixed for every admissible coefficient. + endpoints = np.array([0.0, 1.0], dtype=np.float32) + for s in (SWAY_COEF_MIN, -0.5, 0.5, SWAY_COEF_MAX): + out = sway_sampling(endpoints, s) + assert out[0] == 0.0 + assert out[-1] == 1.0 + # The map is monotone increasing on the admissible interval. + fine = np.linspace(0.0, 1.0, 1001, dtype=np.float32) + mapped = sway_sampling(fine, DEFAULT_SWAY_COEF) + assert np.all(np.diff(mapped) > 0) + + +def test_sway_sampling_rejects_non_monotone_coefficients(): + u = np.linspace(0.0, 1.0, 5, dtype=np.float32) + with pytest.raises(ValueError): + sway_sampling(u, SWAY_COEF_MIN - 0.5) + with pytest.raises(ValueError): + sway_sampling(u, SWAY_COEF_MAX + 0.5) + + +def test_sway_sampling_works_on_torch_tensors(): + u = torch.linspace(0.0, 1.0, 9) + out = sway_sampling(u, DEFAULT_SWAY_COEF) + assert isinstance(out, torch.Tensor) + assert out.shape == u.shape + np_out = sway_sampling(u.numpy(), DEFAULT_SWAY_COEF) + np.testing.assert_allclose(out.numpy(), np_out, atol=1e-6) + + +def _make_scheduler(num_inference_steps=8, shift=1.0): + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift) + scheduler.set_timesteps(num_inference_steps) + return scheduler + + +def test_apply_sway_sampling_preserves_schedule_invariants(): + scheduler = _make_scheduler(num_inference_steps=8) + orig_sigmas = scheduler.sigmas.clone() + orig_timesteps = scheduler.timesteps.clone() + + apply_sway_sampling(scheduler, sway_coef=DEFAULT_SWAY_COEF) + + # Lengths are unchanged: N timesteps, N+1 sigmas (terminal 0 included). + assert scheduler.sigmas.shape[0] == orig_sigmas.shape[0] + assert scheduler.timesteps.shape[0] == orig_timesteps.shape[0] + # The trajectory endpoints (start noise level, clean 0) are preserved. + assert torch.allclose(scheduler.sigmas[0], orig_sigmas[0]) + assert torch.allclose(scheduler.sigmas[-1], orig_sigmas[-1]) + # Sigmas stay strictly decreasing along the interior. + assert bool(torch.all(scheduler.sigmas[:-1] > scheduler.sigmas[1:])) + # Timesteps remain consistent with the resampled sigmas. + assert torch.allclose(scheduler.timesteps, scheduler.sigmas[:-1] * scheduler.config.num_train_timesteps) + + +def test_apply_sway_sampling_zero_coefficient_is_identity(): + scheduler = _make_scheduler(num_inference_steps=8) + before = scheduler.sigmas.clone() + apply_sway_sampling(scheduler, sway_coef=0.0) + # s = 0 maps every flow step to itself, so the schedule is recovered exactly. + assert torch.allclose(scheduler.sigmas, before, atol=1e-6) + + +def test_apply_sway_sampling_default_biases_steps_toward_start(): + # With s = -1, flow steps are pulled toward the start of the trajectory + # (high noise level), so the midpoint step lands at a higher sigma than it + # does on the uniform schedule. + scheduler = _make_scheduler(num_inference_steps=8) + uniform_midpoint_sigma = scheduler.sigmas[4].clone() + apply_sway_sampling(scheduler, sway_coef=DEFAULT_SWAY_COEF) + assert bool(scheduler.sigmas[4] > uniform_midpoint_sigma)