|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Emit C++ goldens for the S2Mel DiT FRONT END: how x_in is built. |
| 3 | +
|
| 4 | +Upstream `indextts/s2mel/modules/diffusion_transformer.py:243-253`, index-tts |
| 5 | +@4f8792ff120cd3ea470dd511e997a17c86cddd10, under the shipped config |
| 6 | +(`long_skip_connection: true`, `final_layer_type: wavenet`): |
| 7 | +
|
| 8 | + if long_skip_connection: x_res = skip_linear(cat([x_res, x], dim=-1)) |
| 9 | + x = conv1(x_res) # Linear D -> wavenet hidden |
| 10 | + x = x.transpose(1, 2) # [B, H, T] |
| 11 | + t2 = t_embedder2(t) |
| 12 | + x = wavenet(x, x_mask, g=t2.unsqueeze(2)).transpose(1, 2) + res_projection(x_res) |
| 13 | + x = final_layer(x, t1).transpose(1, 2) |
| 14 | + x = conv2(x) # Conv1d H -> in_channels, kernel 1 |
| 15 | +
|
| 16 | +The DiT is constructed for real at reduced dims, so every module here is |
| 17 | +upstream's own; only the SEQUENCE is restated, and it is restated once, next to |
| 18 | +the upstream line numbers it copies. |
| 19 | +
|
| 20 | +Usage: DIT_SRC=<path to indextts/s2mel/modules> python3 \ |
| 21 | + scripts/gen-dit-front-goldens.py --out tests/vllm/models/dit_front_goldens.inc |
| 22 | +""" |
| 23 | + |
| 24 | +from __future__ import annotations |
| 25 | + |
| 26 | +import argparse |
| 27 | +import importlib.util |
| 28 | +import os |
| 29 | +import sys |
| 30 | +import types |
| 31 | +from pathlib import Path |
| 32 | +from types import SimpleNamespace |
| 33 | + |
| 34 | +import torch |
| 35 | + |
| 36 | + |
| 37 | +def rnd(name: str, n: int, scale: float = 1.0) -> list: |
| 38 | + h = 0xCBF29CE484222325 |
| 39 | + for ch in name.encode(): |
| 40 | + h = ((h ^ ch) * 0x100000001B3) & 0xFFFFFFFFFFFFFFFF |
| 41 | + out = [] |
| 42 | + for _ in range(n): |
| 43 | + h = (h + 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF |
| 44 | + z = h |
| 45 | + z = ((z ^ (z >> 30)) * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF |
| 46 | + z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & 0xFFFFFFFFFFFFFFFF |
| 47 | + z ^= z >> 31 |
| 48 | + out.append(((z >> 11) * (1.0 / 9007199254740992.0) * 2.0 - 1.0) * scale) |
| 49 | + return out |
| 50 | + |
| 51 | + |
| 52 | +def tensor(name: str, shape, scale: float = 1.0) -> torch.Tensor: |
| 53 | + n = 1 |
| 54 | + for d in shape: |
| 55 | + n *= d |
| 56 | + return torch.tensor(rnd(name, n, scale), dtype=torch.float64).reshape(shape).float() |
| 57 | + |
| 58 | + |
| 59 | +def load_dit(src: Path): |
| 60 | + sys.path.insert(0, str(src.parents[2])) |
| 61 | + for name in ("munch",): |
| 62 | + if name not in sys.modules: |
| 63 | + stub = types.ModuleType(name) |
| 64 | + stub.Munch = dict |
| 65 | + sys.modules[name] = stub |
| 66 | + spec = importlib.util.spec_from_file_location( |
| 67 | + "indextts.s2mel.modules.diffusion_transformer", src / "diffusion_transformer.py" |
| 68 | + ) |
| 69 | + m = importlib.util.module_from_spec(spec) |
| 70 | + spec.loader.exec_module(m) |
| 71 | + return m |
| 72 | + |
| 73 | + |
| 74 | +def fmt(values) -> str: |
| 75 | + lines, row = [], [] |
| 76 | + for v in values: |
| 77 | + row.append(f"{float(v):.9e}F") |
| 78 | + if len(row) == 6: |
| 79 | + lines.append(" " + ", ".join(row) + ",") |
| 80 | + row = [] |
| 81 | + if row: |
| 82 | + lines.append(" " + ", ".join(row) + ",") |
| 83 | + return "\n".join(lines) |
| 84 | + |
| 85 | + |
| 86 | +# Reduced dims. The shipped model is hidden 512 / wavenet 512 / in_channels 80 / |
| 87 | +# 8 wavenet layers; the RATIOS that matter (skip_linear takes hidden + in_channels, |
| 88 | +# conv2 maps wavenet hidden -> in_channels) are preserved. |
| 89 | +# |
| 90 | +# WN_HIDDEN MUST EQUAL HIDDEN. `final_layer` is built at the wavenet width but is |
| 91 | +# called with `t1`, which the DiT embeds at ITS hidden width, so the wavenet |
| 92 | +# final-layer path only composes when the two are equal. They both happen to be |
| 93 | +# 512 upstream, which hides the coupling; setting them differently here raised |
| 94 | +# `mat1 and mat2 shapes cannot be multiplied (1x8 and 6x12)` from upstream's own |
| 95 | +# module. The C++ port asserts it rather than inheriting a silent coincidence. |
| 96 | +HIDDEN, WN_HIDDEN, IN_CH, HEADS, DEPTH = 8, 8, 4, 2, 1 |
| 97 | +STYLE = 6 |
| 98 | +WN_LAYERS, WN_KERNEL, WN_DILATION, FRAMES = 2, 3, 1, 7 |
| 99 | + |
| 100 | + |
| 101 | +def build_args(): |
| 102 | + dit = SimpleNamespace( |
| 103 | + time_as_token=False, style_as_token=False, uvit_skip_connection=True, |
| 104 | + depth=DEPTH, num_heads=HEADS, hidden_dim=HIDDEN, block_size=128, |
| 105 | + in_channels=IN_CH, content_type="discrete", content_codebook_size=16, |
| 106 | + content_dim=HIDDEN, is_causal=False, final_layer_type="wavenet", |
| 107 | + style_condition=True, class_dropout_prob=0.0, long_skip_connection=True, |
| 108 | + target="mel", f0_condition=False, n_f0_bins=8, content_codebooks=1, |
| 109 | + zero_prompt_speech_token=False, add_resblock_in_transformer=False, |
| 110 | + ) |
| 111 | + wavenet = SimpleNamespace( |
| 112 | + hidden_dim=WN_HIDDEN, num_layers=WN_LAYERS, kernel_size=WN_KERNEL, |
| 113 | + dilation_rate=WN_DILATION, p_dropout=0.0, style_condition=True, |
| 114 | + ) |
| 115 | + style_encoder = SimpleNamespace(dim=STYLE) |
| 116 | + return SimpleNamespace(DiT=dit, wavenet=wavenet, style_encoder=style_encoder) |
| 117 | + |
| 118 | + |
| 119 | +def main() -> int: |
| 120 | + ap = argparse.ArgumentParser() |
| 121 | + ap.add_argument("--out", required=True) |
| 122 | + a = ap.parse_args() |
| 123 | + m = load_dit(Path(os.environ["DIT_SRC"])) |
| 124 | + |
| 125 | + torch.manual_seed(0) |
| 126 | + dit = m.DiT(build_args()).eval() |
| 127 | + |
| 128 | + # Every parameter on the tail comes from the shared stream, so the C++ side |
| 129 | + # rebuilds them without a fixture. |
| 130 | + tail_prefixes = ("cond_projection.", "cond_x_merge_linear.") |
| 131 | + with torch.no_grad(): |
| 132 | + for pname, p in sorted(dit.named_parameters()): |
| 133 | + if pname.startswith(tail_prefixes): |
| 134 | + p.copy_(tensor("front." + pname, list(p.shape), 0.5)) |
| 135 | + |
| 136 | + x = tensor("front.x", [1, IN_CH, FRAMES]) # channel-major, as upstream |
| 137 | + prompt_x = tensor("front.prompt_x", [1, IN_CH, FRAMES]) |
| 138 | + cond = tensor("front.cond", [1, FRAMES, HIDDEN]) |
| 139 | + style = tensor("front.style", [1, STYLE]) |
| 140 | + |
| 141 | + with torch.no_grad(): |
| 142 | + # ---- upstream diffusion_transformer.py:206-226, verbatim order ---- |
| 143 | + cond_p = dit.cond_projection(cond) |
| 144 | + xt = x.transpose(1, 2) |
| 145 | + pt = prompt_x.transpose(1, 2) |
| 146 | + x_in = torch.cat([xt, pt, cond_p], dim=-1) |
| 147 | + x_in = torch.cat([x_in, style[:, None, :].repeat(1, FRAMES, 1)], dim=-1) |
| 148 | + cat864 = x_in.clone() |
| 149 | + merged = dit.cond_x_merge_linear(x_in) |
| 150 | + |
| 151 | + # the CFG unconditional branch: everything past in_channels zeroed |
| 152 | + x_in_u = cat864.clone() |
| 153 | + x_in_u[..., IN_CH:] = x_in_u[..., IN_CH:] * 0 |
| 154 | + merged_u = dit.cond_x_merge_linear(x_in_u) |
| 155 | + # ------------------------------------------------------------------- |
| 156 | + |
| 157 | + names = sorted(n for n, _ in dit.named_parameters() if n.startswith(tail_prefixes)) |
| 158 | + |
| 159 | + body = [ |
| 160 | + "// GENERATED by scripts/gen-dit-front-goldens.py -- do not edit.", |
| 161 | + "// Oracle: diffusion_transformer.py:243-253 (DiT tail), index-tts", |
| 162 | + "// @4f8792ff120cd3ea470dd511e997a17c86cddd10, under the SHIPPED config", |
| 163 | + "// long_skip_connection: true, final_layer_type: wavenet.", |
| 164 | + "#pragma once", |
| 165 | + "", |
| 166 | + "#include <cstdint>", |
| 167 | + "", |
| 168 | + "namespace dit_front_goldens {", |
| 169 | + "", |
| 170 | + f"inline constexpr int64_t kHidden = {HIDDEN};", |
| 171 | + f"inline constexpr int64_t kInChannels = {IN_CH};", |
| 172 | + f"inline constexpr int64_t kStyle = {STYLE};", |
| 173 | + f"inline constexpr int64_t kFrames = {FRAMES};", |
| 174 | + "", |
| 175 | + "inline constexpr const char* kParamNames[] = {", |
| 176 | + ] |
| 177 | + body += [f' "front.{n}",' for n in names] |
| 178 | + body += [ |
| 179 | + "};", |
| 180 | + "", |
| 181 | + "// The 864-wide concatenation before the merge -- [kFrames, 864].", |
| 182 | + "inline constexpr float kCat[] = {", |
| 183 | + fmt(cat864.reshape(-1).tolist()), |
| 184 | + "};", |
| 185 | + "", |
| 186 | + "// cond_x_merge_linear(cat) -- [kFrames, kHidden].", |
| 187 | + "inline constexpr float kMerged[] = {", |
| 188 | + fmt(merged.reshape(-1).tolist()), |
| 189 | + "};", |
| 190 | + "", |
| 191 | + "// The CFG UNCONDITIONAL branch: columns past kInChannels zeroed first.", |
| 192 | + "inline constexpr float kMergedUncond[] = {", |
| 193 | + fmt(merged_u.reshape(-1).tolist()), |
| 194 | + "};", |
| 195 | + "", |
| 196 | + "} // namespace dit_front_goldens", |
| 197 | + "", |
| 198 | + ] |
| 199 | + |
| 200 | + Path(a.out).write_text("\n".join(body)) |
| 201 | + print(f"wrote {a.out}: {len(names)} front params, merged {tuple(merged.shape)}") |
| 202 | + return 0 |
| 203 | + |
| 204 | + |
| 205 | +if __name__ == "__main__": |
| 206 | + raise SystemExit(main()) |
0 commit comments