| """YuE2 Oobleck VAE decoder in MLX (channels-last; weight norm folded by convert.py).""" |
| from __future__ import annotations |
|
|
| import json |
| import math |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
|
|
| FRAME = 1920 |
|
|
|
|
| class Snake(nn.Module): |
| def __init__(self, ch): |
| super().__init__() |
| self.alpha = mx.zeros((ch,)) |
| self.beta = mx.zeros((ch,)) |
|
|
| def __call__(self, x): |
| alpha, beta = mx.exp(self.alpha), mx.exp(self.beta) |
| return x + (1.0 / (beta + 1e-9)) * mx.sin(x * alpha) ** 2 |
|
|
|
|
| class ResidualUnit(nn.Module): |
| def __init__(self, ch, dilation): |
| super().__init__() |
| self.layers = [Snake(ch), nn.Conv1d(ch, ch, 7, dilation=dilation, padding=3 * dilation), |
| Snake(ch), nn.Conv1d(ch, ch, 1)] |
|
|
| def __call__(self, x): |
| y = x |
| for layer in self.layers: |
| y = layer(y) |
| return x + y |
|
|
|
|
| class DecoderBlock(nn.Module): |
| def __init__(self, cin, cout, stride): |
| super().__init__() |
| self.layers = [Snake(cin), |
| nn.ConvTranspose1d(cin, cout, 2 * stride, stride=stride, padding=math.ceil(stride / 2)), |
| ResidualUnit(cout, 1), ResidualUnit(cout, 3), ResidualUnit(cout, 9)] |
|
|
| def __call__(self, x): |
| for layer in self.layers: |
| x = layer(x) |
| return x |
|
|
|
|
| class OobleckDecoder(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| c, mults, strides = cfg["channels"], [1] + list(cfg["c_mults"]), cfg["strides"] |
| layers = [nn.Conv1d(cfg["latent_dim"], mults[-1] * c, 7, padding=3)] |
| for i in range(len(mults) - 1, 0, -1): |
| layers.append(DecoderBlock(mults[i] * c, mults[i - 1] * c, strides[i - 1])) |
| layers += [Snake(mults[0] * c), nn.Conv1d(mults[0] * c, cfg["out_channels"], 7, padding=3, bias=False)] |
| self.layers = layers |
|
|
| def __call__(self, x): |
| for layer in self.layers: |
| x = layer(x) |
| return x |
|
|
| def decode_tiled(self, z, core=1024, halo=16): |
| """z [T,64] float32 → audio [1920*T-64, 2] float32, exact tile cores, no crossfade.""" |
| frames, total = z.shape[0], FRAME * z.shape[0] - 64 |
| out = [] |
| for start in range(0, frames, core): |
| end = min(frames, start + core) |
| left, right = max(0, start - halo), min(frames, end + halo) |
| tile = self(z[None, left:right])[0] |
| crop = (start - left) * FRAME |
| out.append(tile[crop:crop + min(end * FRAME, total) - start * FRAME]) |
| mx.eval(out[-1]) |
| return mx.concatenate(out) |
|
|
|
|
| def load_vae(path: Path) -> OobleckDecoder: |
| path = Path(path) |
| cfg = json.loads((path / "vae_config.json").read_text()) |
| decoder = OobleckDecoder(cfg["decoder_config"]) |
| decoder.load_weights(str(path / "vae.safetensors")) |
| mx.eval(decoder.parameters()) |
| return decoder |
|
|
