From ea72fced601ec7a8c321e08bf52639ade5fa1735 Mon Sep 17 00:00:00 2001 From: Staroom <31646341+eastmoe@users.noreply.github.com> Date: Sat, 23 May 2026 17:45:58 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9C=AC=E5=9C=B0=E6=9B=BF=E4=BB=A3=E4=BE=9D?= =?UTF-8?q?=E8=B5=96k-diffusion=20=E5=92=8C=20vector-quantize-pytorch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 在 requirements.txt (line 20) 中注释掉 k-diffusion 和 vector-quantize-pytorch，并写明本地替代原因。 新增本地兼容层 sampling.py (line 17)，覆盖 get_sigmas_karras、sample_euler、sample_heun、sample_dpmpp_2m 这几个常见推理 API。 在 descript_quantize3.py (line 11) 标注了本地 RVQ 替代范围：F.embedding 做 indices 到 codebook embedding，归一化平方 L2 最近邻做 encoder-side quantize。 --- .../Flow1dVAE/libs/rvq/descript_quantize3.py | 5 +- songgeneration/k_diffusion/__init__.py | 2 + songgeneration/k_diffusion/sampling.py | 159 ++++++++++++++++++ songgeneration/requirements.txt | 7 +- 4 files changed, 168 insertions(+), 5 deletions(-) create mode 100644 songgeneration/k_diffusion/__init__.py create mode 100644 songgeneration/k_diffusion/sampling.py diff --git a/songgeneration/codeclm/tokenizer/Flow1dVAE/libs/rvq/descript_quantize3.py b/songgeneration/codeclm/tokenizer/Flow1dVAE/libs/rvq/descript_quantize3.py index 946d20e..6c6d2a5 100644 --- a/songgeneration/codeclm/tokenizer/Flow1dVAE/libs/rvq/descript_quantize3.py +++ b/songgeneration/codeclm/tokenizer/Flow1dVAE/libs/rvq/descript_quantize3.py @@ -8,6 +8,9 @@ import torch.nn.functional as F from einops import rearrange from torch.nn.utils import weight_norm +# Local replacement for the small vector-quantize-pytorch surface used here: +# - code indices -> vectors is F.embedding in embed_code() +# - encoder-side quantization is nearest-neighbor lookup in decode_latents() def WNConv1d(*args, **kwargs): return weight_norm(nn.Conv1d(*args, **kwargs)) @@ -86,7 +89,7 @@ class VectorQuantize(nn.Module): encodings = F.normalize(encodings) codebook = F.normalize(codebook) - # Compute euclidean distance with codebook + # Compute squared L2 distance with the normalized codebook. dist = ( encodings.pow(2).sum(1, keepdim=True) - 2 * encodings @ codebook.t() diff --git a/songgeneration/k_diffusion/__init__.py b/songgeneration/k_diffusion/__init__.py new file mode 100644 index 0000000..d9f4007 --- /dev/null +++ b/songgeneration/k_diffusion/__init__.py @@ -0,0 +1,2 @@ +"""Local inference-only compatibility subset for k-diffusion.""" + diff --git a/songgeneration/k_diffusion/sampling.py b/songgeneration/k_diffusion/sampling.py new file mode 100644 index 0000000..9ff4261 --- /dev/null +++ b/songgeneration/k_diffusion/sampling.py @@ -0,0 +1,159 @@ +import math + +import torch + + +def append_zero(x): + return torch.cat([x, x.new_zeros([1])]) + + +def append_dims(x, target_dims): + dims_to_append = target_dims - x.ndim + if dims_to_append < 0: + raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}") + return x[(...,) + (None,) * dims_to_append] + + +def get_sigmas_karras(n, sigma_min, sigma_max, rho=7.0, device="cpu"): + """Construct the Karras noise schedule used by k-diffusion samplers.""" + ramp = torch.linspace(0, 1, n, device=device) + min_inv_rho = sigma_min ** (1 / rho) + max_inv_rho = sigma_max ** (1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + return append_zero(sigmas) + + +def to_d(x, sigma, denoised): + """Convert denoiser output to the ODE derivative expected by these samplers.""" + return (x - denoised) / append_dims(sigma, x.ndim) + + +def _default_extra_args(extra_args): + return {} if extra_args is None else extra_args + + +def _call_callback(callback, x, i, sigma, sigma_hat, denoised): + if callback is not None: + callback( + { + "x": x, + "i": i, + "sigma": sigma, + "sigma_hat": sigma_hat, + "denoised": denoised, + } + ) + + +@torch.no_grad() +def sample_euler( + model, + x, + sigmas, + extra_args=None, + callback=None, + disable=None, + s_churn=0.0, + s_tmin=0.0, + s_tmax=float("inf"), + s_noise=1.0, +): + """Minimal Euler sampler compatible with k_diffusion.sampling.sample_euler.""" + extra_args = _default_extra_args(extra_args) + s_in = x.new_ones([x.shape[0]]) + + for i in range(len(sigmas) - 1): + gamma = ( + min(s_churn / (len(sigmas) - 1), math.sqrt(2) - 1) + if s_tmin <= sigmas[i] <= s_tmax + else 0.0 + ) + sigma_hat = sigmas[i] * (gamma + 1) + if gamma > 0: + x = x + ( + torch.randn_like(x) * s_noise * (sigma_hat**2 - sigmas[i] ** 2).sqrt() + ) + + denoised = model(x, sigma_hat * s_in, **extra_args) + d = to_d(x, sigma_hat, denoised) + _call_callback(callback, x, i, sigmas[i], sigma_hat, denoised) + x = x + d * (sigmas[i + 1] - sigma_hat) + + return x + + +@torch.no_grad() +def sample_heun( + model, + x, + sigmas, + extra_args=None, + callback=None, + disable=None, + s_churn=0.0, + s_tmin=0.0, + s_tmax=float("inf"), + s_noise=1.0, +): + """Minimal Heun sampler compatible with k_diffusion.sampling.sample_heun.""" + extra_args = _default_extra_args(extra_args) + s_in = x.new_ones([x.shape[0]]) + + for i in range(len(sigmas) - 1): + gamma = ( + min(s_churn / (len(sigmas) - 1), math.sqrt(2) - 1) + if s_tmin <= sigmas[i] <= s_tmax + else 0.0 + ) + sigma_hat = sigmas[i] * (gamma + 1) + if gamma > 0: + x = x + ( + torch.randn_like(x) * s_noise * (sigma_hat**2 - sigmas[i] ** 2).sqrt() + ) + + denoised = model(x, sigma_hat * s_in, **extra_args) + d = to_d(x, sigma_hat, denoised) + _call_callback(callback, x, i, sigmas[i], sigma_hat, denoised) + dt = sigmas[i + 1] - sigma_hat + + if sigmas[i + 1] == 0: + x = x + d * dt + else: + x_2 = x + d * dt + denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args) + d_2 = to_d(x_2, sigmas[i + 1], denoised_2) + x = x + (d + d_2) * 0.5 * dt + + return x + + +@torch.no_grad() +def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None): + """Minimal DPM++ 2M sampler compatible with k_diffusion.sampling.sample_dpmpp_2m.""" + extra_args = _default_extra_args(extra_args) + s_in = x.new_ones([x.shape[0]]) + old_denoised = None + h_last = None + + for i in range(len(sigmas) - 1): + denoised = model(x, sigmas[i] * s_in, **extra_args) + _call_callback(callback, x, i, sigmas[i], sigmas[i], denoised) + + if sigmas[i + 1] == 0: + x = denoised + else: + t = -sigmas[i].log() + t_next = -sigmas[i + 1].log() + h = t_next - t + + if old_denoised is None or h_last is None: + denoised_d = denoised + else: + r = h_last / h + denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised + + x = (sigmas[i + 1] / sigmas[i]) * x - (-h).expm1() * denoised_d + old_denoised = denoised + h_last = h + + return x diff --git a/songgeneration/requirements.txt b/songgeneration/requirements.txt index fe07ec3..0c0a8ca 100644 --- a/songgeneration/requirements.txt +++ b/songgeneration/requirements.txt @@ -11,17 +11,16 @@ # alias-free-torch>=0.0.6 # replaced by local alias_free_torch shim - # einops-exts>=0.0.4 # replaced by local einops_exts.py helpers backed by einops - # julius>=0.2.7 # replaced by torchaudio.functional.resample shim +# k-diffusion>=0.1.1,<0.2 # replaced by local k_diffusion.sampling inference subset; current Flow1dVAE inference uses its CFM Euler solver +# vector-quantize-pytorch>=1.22.17 # replaced by Flow1dVAE/libs/rvq/descript_quantize3.py local RVQ implementation + diffusers>=0.37.1,<0.38 huggingface-hub>=0.34.0,<2.0 -k-diffusion>=0.1.1,<0.2 lameenc>=1.8.1 librosa>=0.11.0 omegaconf>=2.2.0,<3.0 openunmix>=1.3.0 packaging>=24.0 -vector-quantize-pytorch>=1.22.17 x-transformers>=2.3.25