File size: 7,724 Bytes
5eee449 | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """Host-side chain between encoder.axmodel and decoder.axmodel (pure numpy).
Replicates, in numpy, the torch host chain of export/export_onnx.py:host_expand
(which itself mirrors origin/runtime/models.py:546-558):
w = exp(logw) * x_mask * length_scale
w_ceil = ceil(w); T' = max(sum(w_ceil), 1)
attn = generate_path(w_ceil, mask) # pure integer logic
m_p' = attn @ m_p^T ; logs_p' = attn @ logs_p^T
z_p = m_p' + randn * exp(logs_p') * variation (seed-controlled, host)
Also contains the fixed-shape plumbing: encoder T=256 zero padding and decoder
Tp=512 chunking with overlap crossfade (EXPORT_NOTES §5.3 recommends
>=64-frame overlap between chunks; single-chunk T'<=512 is the path validated
by SIMULATE).
"""
from __future__ import annotations
import numpy as np
SAMPLE_RATE = 24000
HOP_LENGTH = 256
HIDDEN_CHANNELS = 192
ENCODER_T = 256
DECODER_TP = 512
DECODER_OVERLAP = 64 # frames of overlap between decoder chunks (EXPORT_NOTES §5.3)
# ---------------------------------------------------------------------------
# Duration rounding + alignment (generate_path)
# ---------------------------------------------------------------------------
def expand_priors(
logw: np.ndarray,
m_p: np.ndarray,
logs_p: np.ndarray,
x_len: int,
length_scale: float,
) -> tuple[np.ndarray, np.ndarray, int]:
"""Duration rounding + generate_path + prior expansion.
Args:
logw: [1, 1, T] log-durations (encoder output, static T)
m_p: [1, C, T] prior means
logs_p: [1, C, T] prior log-stds
x_len: number of valid (unpadded) token frames
length_scale: 1.0 / speed
Returns:
(m_p_e [T', C], logs_p_e [T', C], T')
"""
t_total = logw.shape[-1]
if not 1 <= x_len <= t_total:
raise ValueError(f"x_len {x_len} out of range [1, {t_total}]")
x_mask = (np.arange(t_total) < x_len).astype(np.float32)
w = np.exp(logw[0, 0].astype(np.float32)) * x_mask * np.float32(length_scale)
w_ceil = np.ceil(w).astype(np.int64)
y_len = max(int(w_ceil.sum()), 1)
# generate_path (origin/runtime/commons.py): with the masked durations the
# alignment reduces to pure integer interval logic — frame t of the output
# copies input frame i for cum[i-1] <= t < cum[i]. O(T) loop, no matmul
# of the [T', T] path needed; we still build attn explicitly for the prior
# expansion to stay a literal port (T=256, T'<=~2000: cheap).
cum = np.cumsum(w_ceil)
attn = np.zeros((y_len, t_total), dtype=np.float32)
lo = 0
for i in range(t_total):
hi = min(int(cum[i]), y_len)
if hi > lo:
attn[lo:hi, i] = 1.0
lo = hi
m_p_e = attn @ m_p[0].T.astype(np.float32)
logs_p_e = attn @ logs_p[0].T.astype(np.float32)
return m_p_e, logs_p_e, y_len
def inject_noise(
m_p_e: np.ndarray,
logs_p_e: np.ndarray,
variation: float,
seed: int,
) -> np.ndarray:
"""z_p = m_p' + randn * exp(logs_p') * variation, seed-controlled.
NOTE: numpy PCG64 is used, not torch's MT19937 — a given `seed` is
deterministic and reproducible within this SDK but is NOT bit-identical to
the PyTorch reference (origin/inference.py).
"""
rng = np.random.default_rng(seed)
noise = rng.standard_normal(m_p_e.shape, dtype=np.float32)
z_p = m_p_e + noise * np.exp(logs_p_e.astype(np.float32)) * np.float32(variation)
return z_p.T.copy() # [C, T']
# ---------------------------------------------------------------------------
# Decoder chunk composition (Tp=512 static shape, overlap crossfade)
# ---------------------------------------------------------------------------
def decoder_chunk_starts(t_prime: int, chunk: int = DECODER_TP,
overlap: int = DECODER_OVERLAP) -> list[int]:
"""Frame offsets of decoder chunks covering [0, t_prime) with >=overlap."""
if t_prime <= chunk:
return [0]
stride = chunk - overlap
starts = list(range(0, t_prime - chunk + 1, stride))
last = t_prime - chunk
if starts[-1] < last:
starts.append(last)
return starts
def decode_waveform(
z_p: np.ndarray,
run_decoder,
chunk: int = DECODER_TP,
overlap: int = DECODER_OVERLAP,
) -> np.ndarray:
"""Run the fixed-shape decoder over z_p and stitch chunks.
Args:
z_p: [C, T'] float32 prior with noise injected.
run_decoder: callable taking z_p_chunk [C, chunk] float32 (already
right-padded with zeros) and returning wav [chunk*HOP] float32.
chunk/overlap: see DECODER_TP / DECODER_OVERLAP.
Returns:
wav [T'*HOP] float32 (tail trimmed; chunk overlaps crossfaded).
"""
t_prime = z_p.shape[1]
hop = HOP_LENGTH
starts = decoder_chunk_starts(t_prime, chunk, overlap)
out: np.ndarray | None = None
prev_end_frame = 0
for start in starts:
take = min(chunk, t_prime - start)
z_chunk = np.zeros((z_p.shape[0], chunk), dtype=np.float32)
z_chunk[:, :take] = z_p[:, start:start + take]
wav = np.asarray(run_decoder(z_chunk), dtype=np.float32)
if out is None:
out = wav
else:
ov_frames = prev_end_frame - start
ov = ov_frames * hop
fade_in = np.linspace(0.0, 1.0, ov, endpoint=True, dtype=np.float32)
out = np.concatenate(
[out[:-ov], out[-ov:] * (1.0 - fade_in) + wav[:ov] * fade_in, wav[ov:]]
)
prev_end_frame = start + chunk
assert out is not None
return out[: t_prime * hop]
# ---------------------------------------------------------------------------
# Waveform post-processing (origin/inference.py)
# ---------------------------------------------------------------------------
def edge_fade(waveform: np.ndarray, sample_rate: int = SAMPLE_RATE,
milliseconds: float = 5.0) -> np.ndarray:
frames = min(round(sample_rate * milliseconds / 1000.0), waveform.size // 2)
if frames <= 0:
return waveform
output = waveform.copy()
ramp = np.linspace(0.0, 1.0, frames, endpoint=True, dtype=np.float32)
output[:frames] *= ramp
output[-frames:] *= ramp[::-1]
return output
# ---------------------------------------------------------------------------
# Text chunking (origin/inference.py: split_text / boundary_pause_seconds)
# ---------------------------------------------------------------------------
def split_text(text: str, limit: int = 280) -> list[str]:
import re
normalized = " ".join(text.split())
sentences = [
part.strip()
for part in re.split(r"(?<=[.!?;:])\s+", normalized)
if part.strip()
]
chunks: list[str] = []
for sentence in sentences or [normalized]:
while len(sentence) > limit:
search = sentence[: limit + 1]
punctuation = max(search.rfind(mark) for mark in (",", ";", ":"))
split_at = (
punctuation + 1
if punctuation >= limit // 2
else sentence.rfind(" ", 0, limit + 1)
)
if split_at < limit // 2:
split_at = limit
chunks.append(sentence[:split_at].strip())
sentence = sentence[split_at:].strip()
if sentence:
chunks.append(sentence)
return chunks
def boundary_pause_seconds(chunk: str) -> float:
ending = chunk.rstrip()[-1:] if chunk.strip() else ""
return {
"?": 0.28,
"!": 0.24,
".": 0.22,
";": 0.16,
":": 0.13,
",": 0.09,
}.get(ending, 0.08)
def seconds_to_samples(seconds: float) -> int:
return round(seconds * SAMPLE_RATE)
|