File size: 3,909 Bytes
2577656 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | # Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87.
# Licensed under Apache-2.0; see LICENSE.
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from numbers import Integral
import torch
from .protocol import CODEC_OFFSET, CODEC_SIZE, CONTEXT, MUSIC_END, chunk_ranges
@dataclass
class Chunk:
ar_tokens: list[int]
noise: torch.Tensor
class YuE2PrefixKVCache:
"""Holds a chunk's token-prefix keys and values; later acoustic calls attend to the prefix without extending it."""
def __init__(self):
self.keys, self.values = [], []
def get_seq_length(self):
return self.keys[0].shape[1] if self.keys else 0
def update(self, key, value, layer_idx):
if layer_idx == len(self.keys):
self.keys.append(key)
self.values.append(value)
return key, value
return torch.cat((self.keys[layer_idx], key), dim=1), torch.cat((self.values[layer_idx], value), dim=1)
def _integers(values, name):
result = list(values)
if not result or any(isinstance(v, bool) or not isinstance(v, Integral) for v in result):
raise ValueError(f"{name} must be a nonempty sequence of integer token IDs")
return [int(v) for v in result]
def song_chunks(prefix, codec, seed, context=CONTEXT):
"""Draw the whole song's CPU FP32 noise once, then split it at the release's chunk boundaries."""
prefix = _integers(prefix, "prefix")
codec = _integers(codec, "codec")
if min(prefix) < 0 or min(codec) < 0 or max(codec) >= CODEC_SIZE:
raise ValueError("Token IDs are outside their allowed vocabulary")
if isinstance(context, bool) or not isinstance(context, Integral) or not 1 <= context <= CONTEXT:
raise ValueError(f"context must be an integer in 1..{CONTEXT}")
ranges = chunk_ranges(len(codec), len(prefix), int(context))
generator = torch.Generator(device="cpu").manual_seed(int(seed))
noise = torch.randn((len(codec), 64), dtype=torch.float32, device="cpu", generator=generator)
return [
Chunk(prefix + [value + CODEC_OFFSET for value in codec[a:b]] + [MUSIC_END], noise[a:b]) for a, b in ranges
]
@torch.inference_mode()
def solve_midpoint(
transformer,
kv_cache,
noise,
device,
steps=32,
cancelled: Callable[[], bool] | None = None,
on_progress: Callable[[int, int], None] | None = None,
):
"""Integrate the flow from t=1 (noise) to t=0 with the midpoint method; returns CPU FP32 [frames, 64] latents."""
if isinstance(steps, bool) or not isinstance(steps, Integral) or steps < 1:
raise ValueError("steps must be a positive integer")
if not torch.isfinite(noise).all():
raise ValueError("Acoustic noise contains non-finite values")
def velocity(state, t):
# The model takes time in logit space; logit(1) is clamped to 20.
raw = torch.logit(torch.tensor(t, dtype=torch.float64, device="cpu")).clamp(-20, 20).item()
return transformer(latents=state[None], timestep=raw, kv_cache=kv_cache).sample[0]
state = noise.to(device=device, dtype=transformer.dtype)
dt = 1.0 / steps
for step in range(steps):
if cancelled is not None and cancelled():
raise InterruptedError("Cancelled during acoustic flow matching")
t = 1.0 - step * dt
mid = state - velocity(state, t) * (dt / 2)
if cancelled is not None and cancelled():
raise InterruptedError("Cancelled during acoustic flow matching")
state = state - velocity(mid, t - dt / 2) * dt
if on_progress is not None:
on_progress(step + 1, int(steps))
result = state.float().cpu()
if not torch.isfinite(result).all():
raise FloatingPointError("Acoustic flow matching produced non-finite latents")
return result
|