|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Raon 1B DiT per-stage diff (#387-adj) — localize the C++ DiT divergence. |
| 3 | +
|
| 4 | +The 1B converts+loads+runs but emits a wrong mel (non-speech). Python reference |
| 5 | +is perfect and the GGUF is provably correct, so the bug is in our C++ DiT. This |
| 6 | +kernel captures ONE reference DiT forward (torch, CPU) — the input-embed output |
| 7 | +(hidden), the timestep embedding, every transformer-block output, and the final |
| 8 | +velocity — then injects the SAME hidden + t_emb into our C++ via |
| 9 | +CRISPASR_F5_DIT_PROBE and diffs each stage. The first block whose cosine drops |
| 10 | +is where our port diverges (rope/attn/norm/ffn/adaln all live inside the block). |
| 11 | +
|
| 12 | +POSITIVE CONTROL (a green wall must not be the comparator failing to fire): |
| 13 | + - identity cos(ref_b0, ref_b0) == 1 |
| 14 | + - sensitivity cos(ref_b0, ref_b1) < 0.999 (comparator can report MISMATCH) |
| 15 | + - injection the C++ MUST at least reproduce block 0 from the injected hidden; |
| 16 | + if cpp_block_0 already ~0 or NaN the harness (not the model) is broken. |
| 17 | +
|
| 18 | +CPU torch for the reference (P100 lacks torch stft kernels); C++ runs on GPU |
| 19 | +(CUDA) to match the failing roundtrip config. RAON_SIZE=1B. chr1s4 datasets. |
| 20 | +""" |
| 21 | +import json, os, subprocess, sys, time |
| 22 | +from pathlib import Path |
| 23 | +import numpy as np |
| 24 | + |
| 25 | +WORK = Path("/kaggle/working"); TMP = Path("/kaggle/temp"); TMP.mkdir(exist_ok=True) |
| 26 | +PROBE = TMP / "probe"; PROBE.mkdir(exist_ok=True) |
| 27 | +SIZE = os.environ.get("RAON_SIZE", "1B") |
| 28 | +REPO = f"KRAFTON/Raon-OpenTTS-{SIZE}" |
| 29 | +CKPT_FILE = {"0.3B": "model_225000.pt", "1B": "model_520000.pt"}[SIZE] |
| 30 | +CRISPASR_URL = "https://github.com/CrispStrobe/CrispASR.git" |
| 31 | +CRISPASR_REF = os.environ.get("CRISPASR_REF", "feat/raon-opentts-1b") |
| 32 | +RAON_URL = "https://github.com/krafton-ai/Raon-OpenTTS.git" |
| 33 | +CLONE = TMP / "CrispASR"; RAON = TMP / "Raon-OpenTTS" |
| 34 | + |
| 35 | + |
| 36 | +def sh(cmd, **kw): return subprocess.run(cmd, shell=True, capture_output=True, text=True, **kw) |
| 37 | +def step(name, **kv): print(f"[{time.strftime('%H:%M:%S')}] {name} " + json.dumps(kv), flush=True) |
| 38 | + |
| 39 | + |
| 40 | +for url, dst, ref in ((CRISPASR_URL, CLONE, CRISPASR_REF), (RAON_URL, RAON, None)): |
| 41 | + if not dst.exists(): |
| 42 | + for _ in range(4): |
| 43 | + cmd = ["git", "clone", "--depth", "1"] + (["--branch", ref] if ref else []) + [url, str(dst)] |
| 44 | + if subprocess.run(cmd).returncode == 0: |
| 45 | + break |
| 46 | + time.sleep(15) |
| 47 | +for _ in range(4): |
| 48 | + if subprocess.run(["git", "submodule", "update", "--init", "--recursive", "ggml", "third_party/c2pa-audio"], |
| 49 | + cwd=str(CLONE)).returncode == 0 or (CLONE / "ggml" / "CMakeLists.txt").exists(): |
| 50 | + break |
| 51 | + time.sleep(15) |
| 52 | +sys.path.insert(0, str(CLONE / "tools" / "kaggle")) |
| 53 | +import kaggle_harness as kh # noqa: E402 |
| 54 | +kh.init_progress() |
| 55 | +HF_TOKEN = kh.resolve_hf_token() |
| 56 | +subprocess.run([sys.executable, "-m", "pip", "install", "-q", "x_transformers", "torchdiffeq", "ema_pytorch", |
| 57 | + "loguru", "einops", "jieba", "pypinyin", "hydra-core", "omegaconf", "vocos", "pyyaml", |
| 58 | + "soundfile", "huggingface_hub"], check=False) |
| 59 | +subprocess.run([sys.executable, "-m", "pip", "install", "-q", "--no-deps", "gguf"], check=False) |
| 60 | +sys.path.insert(0, str(RAON / "src")) |
| 61 | +from huggingface_hub import hf_hub_download # noqa: E402 |
| 62 | +import torch, yaml # noqa: E402 |
| 63 | + |
| 64 | +MODELS = TMP / "models"; MODELS.mkdir(exist_ok=True) |
| 65 | +def dl(repo, f, sub=""): |
| 66 | + d = MODELS / sub if sub else MODELS; d.mkdir(parents=True, exist_ok=True) |
| 67 | + return hf_hub_download(repo, f, local_dir=str(d), token=HF_TOKEN or None) |
| 68 | + |
| 69 | + |
| 70 | +# ── build crispasr (GPU, to match the failing config) ────────────────────── |
| 71 | +kh.install_build_toolchain() |
| 72 | +arch = kh.detect_cuda_arch() |
| 73 | +flags = ["-DCMAKE_BUILD_TYPE=Release"] + kh.cuda_build_flags(arch) + kh.cache_and_link_flags() |
| 74 | +r = sh(f"cd {CLONE} && cmake -G Ninja -B build " + " ".join(flags), timeout=1200) |
| 75 | +if r.returncode != 0: |
| 76 | + step("cmake_FAIL", err=r.stderr[-1500:]); sys.exit(1) |
| 77 | +with kh.build_heartbeat("build", interval_s=30): |
| 78 | + kh.sh_with_progress(f"cmake --build build -j{kh.safe_build_jobs(gpu=True)} --target crispasr-cli", cwd=str(CLONE)) |
| 79 | +CLI = CLONE / "build" / "bin" / "crispasr" |
| 80 | +if not CLI.exists(): |
| 81 | + c = [p for p in (CLONE / "build").rglob("crispasr") if p.is_file() and os.access(p, os.X_OK)] |
| 82 | + CLI = c[0] if c else None |
| 83 | +if not CLI: |
| 84 | + step("no_binary"); sys.exit(1) |
| 85 | +os.environ["LD_LIBRARY_PATH"] = str(CLI.parent) + ":" + os.environ.get("LD_LIBRARY_PATH", "") |
| 86 | +gguf = hf_hub_download(f"cstr/raon-opentts-{SIZE.lower()}-GGUF", f"raon-opentts-{SIZE.lower()}-f16.gguf", |
| 87 | + local_dir=str(MODELS), token=HF_TOKEN or None) |
| 88 | +ref_wav = CLONE / "samples" / "jfk.wav" |
| 89 | +step("built", cli=str(CLI), gguf=os.path.basename(gguf)) |
| 90 | + |
| 91 | +# ── reference capture: one DiT forward (torch, CPU, b=1, no cfg) ─────────── |
| 92 | +ckpt = dl(REPO, CKPT_FILE, SIZE); cfg = dl(REPO, "config.yaml", SIZE); vocab = dl(REPO, "vocab.txt", SIZE) |
| 93 | +from f5_tts.model.backbones.dit import DiT # noqa: E402 |
| 94 | +from f5_tts.model.utils import get_tokenizer, list_str_to_idx # noqa: E402 |
| 95 | +conf = yaml.safe_load(open(cfg)); a = conf["model"]["arch"]; mspec = conf["model"]["mel_spec"] |
| 96 | +n_mel = mspec["n_mel_channels"] |
| 97 | +sd = torch.load(ckpt, map_location="cpu", weights_only=True)["ema_model_state_dict"] |
| 98 | +sd = {k.replace("ema_model.transformer.", ""): v for k, v in sd.items() |
| 99 | + if k.startswith("ema_model.transformer.")} |
| 100 | +rows = int(sd["text_embed.text_embed.weight"].shape[0]) |
| 101 | +vmap, _ = get_tokenizer(vocab, "custom") |
| 102 | +vmap = {c: i for c, i in vmap.items() if i < rows - 1} |
| 103 | +dit = DiT(**{k: a[k] for k in a if k != "name"}, mel_dim=n_mel, text_num_embeds=rows - 1) |
| 104 | +missing, unexpected = dit.load_state_dict(sd, strict=False) |
| 105 | +dit.eval() |
| 106 | +step("dit_loaded", missing=len(missing), unexpected=len(unexpected), dim=a["dim"], depth=a["depth"], heads=a["heads"]) |
| 107 | + |
| 108 | +T = 200 |
| 109 | +torch.manual_seed(387) |
| 110 | +x = torch.randn(1, T, n_mel) * 0.3 |
| 111 | +cond = torch.randn(1, T, n_mel) * 0.3 |
| 112 | +gen_text = "the quick brown fox jumps over the lazy dog" |
| 113 | +text_ids = list_str_to_idx([list(gen_text)], vmap) # (1, nt) |
| 114 | +time_t = torch.tensor([0.5]) |
| 115 | + |
| 116 | +dim = a["dim"]; text_dim = a["text_dim"] |
| 117 | +cap = {}; blk_out = {} |
| 118 | +h1 = dit.input_embed.register_forward_hook(lambda m, i, o: cap.__setitem__("hidden", o.detach())) |
| 119 | +h2 = dit.time_embed.register_forward_hook(lambda m, i, o: cap.__setitem__("temb", o.detach())) |
| 120 | +h4 = dit.text_embed.register_forward_hook(lambda m, i, o: cap.__setitem__("text_embed", o.detach())) |
| 121 | +h3 = dit.proj_out.register_forward_hook(lambda m, i, o: cap.__setitem__("velocity", o.detach())) |
| 122 | +bh = [dit.transformer_blocks[k].register_forward_hook( |
| 123 | + (lambda kk: (lambda m, i, o: blk_out.__setitem__(kk, o.detach())))(k)) |
| 124 | + for k in range(len(dit.transformer_blocks))] |
| 125 | +with torch.no_grad(): |
| 126 | + _ = dit(x, cond, text_ids, time_t, drop_audio_cond=False, drop_text=False, cfg_infer=False) |
| 127 | +for h in [h1, h2, h4, h3] + bh: |
| 128 | + h.remove() |
| 129 | + |
| 130 | +hidden = cap["hidden"][0].contiguous().numpy().astype(np.float32) # (T, dim) |
| 131 | +temb = cap["temb"].reshape(-1).numpy().astype(np.float32) # (dim,) |
| 132 | +text_embed = cap["text_embed"][0].contiguous().numpy().astype(np.float32) # (T, text_dim) |
| 133 | +velocity = cap["velocity"][0].contiguous().numpy().astype(np.float32) # (T, mel) |
| 134 | +depth = len(blk_out) |
| 135 | +ref_blocks = [blk_out[k][0].contiguous().numpy().astype(np.float32) for k in range(depth)] |
| 136 | + |
| 137 | +# ── UNCOND arm capture (CFG null: drop_audio_cond + drop_text) ───────────── |
| 138 | +capu = {} |
| 139 | +u1 = dit.input_embed.register_forward_hook(lambda m, i, o: capu.__setitem__("hidden", o.detach())) |
| 140 | +u4 = dit.text_embed.register_forward_hook(lambda m, i, o: capu.__setitem__("text_embed", o.detach())) |
| 141 | +u3 = dit.proj_out.register_forward_hook(lambda m, i, o: capu.__setitem__("velocity", o.detach())) |
| 142 | +dit.clear_cache() |
| 143 | +with torch.no_grad(): |
| 144 | + _ = dit(x, cond, text_ids, time_t, drop_audio_cond=True, drop_text=True, cfg_infer=False) |
| 145 | +for h in [u1, u4, u3]: |
| 146 | + h.remove() |
| 147 | +u_text_embed = capu["text_embed"][0].contiguous().numpy().astype(np.float32) |
| 148 | +u_hidden = capu["hidden"][0].contiguous().numpy().astype(np.float32) |
| 149 | +u_velocity = capu["velocity"][0].contiguous().numpy().astype(np.float32) |
| 150 | +tok = text_ids[0].numpy().astype(np.int32) |
| 151 | +(PROBE / "shape.txt").write_text(f"{T} {tok.size}") |
| 152 | +(PROBE / "t.txt").write_text(str(float(time_t.item()))) |
| 153 | +x[0].contiguous().numpy().astype(np.float32).tofile(PROBE / "x.bin") |
| 154 | +cond[0].contiguous().numpy().astype(np.float32).tofile(PROBE / "cond.bin") |
| 155 | +tok.tofile(PROBE / "tokens.bin") |
| 156 | +hidden.tofile(PROBE / "hidden.bin"); temb.tofile(PROBE / "temb.bin") |
| 157 | +step("captured", T=T, nt=int(tok.size), dim=dim, depth=depth, |
| 158 | + hidden_std=float(hidden.std()), velocity_std=float(velocity.std())) |
| 159 | + |
| 160 | +def run(mode): |
| 161 | + env = dict(os.environ); env[mode] = str(PROBE) |
| 162 | + r = sh(f"{CLI} --backend raon-1b -m {gguf} --voice {ref_wav} --ref-text x --tts probe " |
| 163 | + f"--tts-output {WORK/'probe.wav'} -t 4 --i-have-rights -v", env=env, timeout=1200) |
| 164 | + print(f"[{mode}]", r.stdout[-500:], r.stderr[-1500:], flush=True) |
| 165 | + |
| 166 | +run("CRISPASR_F5_INPUT_PROBE") # -> cpp_temb / cpp_text_embed / cpp_hidden |
| 167 | +run("CRISPASR_F5_DIT_PROBE") # -> cpp_velocity / cpp_block_<k> (injected hidden+temb) |
| 168 | + |
| 169 | +def load(name): |
| 170 | + p = PROBE / name |
| 171 | + return np.fromfile(p, dtype=np.float32) if p.exists() else None |
| 172 | + |
| 173 | +def compare(ref, cpp): |
| 174 | + if cpp is None: |
| 175 | + return None |
| 176 | + u = np.asarray(ref).reshape(-1).astype(np.float64); v = np.asarray(cpp).reshape(-1).astype(np.float64) |
| 177 | + if u.size != v.size: |
| 178 | + return {"size_mismatch": [int(u.size), int(v.size)]} |
| 179 | + nu, nv = np.linalg.norm(u), np.linalg.norm(v) |
| 180 | + cosv = float(u.dot(v) / (nu * nv)) if nu > 0 and nv > 0 else None |
| 181 | + return {"cos": None if cosv is None else round(cosv, 5), |
| 182 | + "norm_ratio": round(float(nv / nu), 4) if nu > 0 else None, # ||cpp|| / ||ref|| (1.0 = same scale) |
| 183 | + "maxabs_ratio": round(float(np.abs(v).max() / (np.abs(u).max() + 1e-12)), 4)} |
| 184 | + |
| 185 | +# positive control on magnitude too: a 1.7x-scaled ref must read norm_ratio≈1.7, cos≈1 |
| 186 | +pc = compare(hidden, (hidden * 1.7).reshape(-1)) |
| 187 | +result = {"size": SIZE, "T": T, "depth": depth, |
| 188 | + "positive_control_scale1.7": pc, |
| 189 | + "cond_input_path": {"temb": compare(temb, load("cpp_temb.bin")), |
| 190 | + "text_embed": compare(text_embed, load("cpp_text_embed.bin")), |
| 191 | + "hidden": compare(hidden, load("cpp_hidden.bin"))}, |
| 192 | + "UNCOND_arm": {"text_embed": compare(u_text_embed, load("cpp_text_embed_uncond.bin")), |
| 193 | + "hidden": compare(u_hidden, load("cpp_hidden_uncond.bin")), |
| 194 | + "velocity": compare(u_velocity, load("cpp_velocity_uncond.bin"))}, |
| 195 | + "cond_velocity": compare(velocity, load("cpp_velocity.bin")), |
| 196 | + "blocks": [{"k": k, **(compare(ref_blocks[k], load(f"cpp_block_{k}.bin")) or {})} for k in range(depth)]} |
| 197 | +(WORK / "raon_dit_diff.json").write_text(json.dumps(result, indent=2)) |
| 198 | +print(json.dumps(result, indent=2), flush=True) |
| 199 | +nr = [b.get("norm_ratio") for b in result["blocks"] if b.get("norm_ratio")] |
| 200 | +step("DONE", cond_velocity=result["cond_velocity"], uncond=result["UNCOND_arm"]["velocity"], block0_nr=result["blocks"][0].get("norm_ratio"), |
| 201 | + blocklast_nr=result["blocks"][-1].get("norm_ratio"), |
| 202 | + norm_ratio_span=[min(nr), max(nr)] if nr else None) |
0 commit comments