本地替代依赖k-diffusion 和 vector-quantize-pytorch

在 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。

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()

songgeneration/k_diffusion/__init__.py

@@ -0,0 +1,2 @@
+"""Local inference-only compatibility subset for k-diffusion."""
+

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

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