~/satyajit

AuK: 6.12 GB, an empty Hub config, and a parameter count that only balances in FP32

mdjsonmcp

2026-09-09 · 28 min · speech · tts · audio · diffusion · open-weights · explainer

AuK is Tencent's open-source speech foundation model — one model for zero-shot TTS, instruction-controlled TTS, content and acoustic editing, emotion/timbre/accent editing, and speech enhancement and separation, all through natural-language instructions. Two links landed on my desk for this: the Hugging Face repo and a PapersWithCode page pointing at arXiv 2609.08936. They're the same artifact — "AuK Technical Report: An Open-Source Foundational Model for Speech Generation and Editing" (Ma, Niu, et al., Tencent, submitted 8 Sep 2026) is the paper behind the weights, and both the model card and the paper cite each other. One article, then, covering both.

PaperarXiv 2609.08936
Weightshuggingface.co/tencent/AuK (base) · tencent/AuK-Flash (4-step distilled, separate repo)
Codegithub.com/Tencent-Hunyuan/AuK — MIT
Repo (base)8 files, usedStorage 6,762,642,715 bytes; auk_base.safetensors 6,122,209,092 B; vae.safetensors 637,322,604 B; not gated; created 2026-08-18, last modified 2026-09-10
BackboneFlux2Edit — 10 dual-stream MMDiT blocks + 20 single-stream DiT blocks, dim 1536, 24×64 heads, ff×2
Semantic encoderQwen2.5-Omni-3B Thinker, vision tower deleted at load — 4,034,780,160 parameters resident, frozen, downloaded separately
Demohuggingface.co/spaces/tencent/AuK — Gradio, ZeroGPU A10G
tencent/AuK@61f0e44 · snapshot 2026-09-09
repo size
6.76 GB
task
text-to-speech
license
mit
safetensors
2 shards
largest file
6.12 GB
files
8
downloads
30
likes
17
audiospeechtext-to-speechzero-shot-ttsvoice-cloningspeech-generationspeech-editingspeech-enhancement

repo last modified 2026-09-09

The Hub's own metadata for this repo is unusually thin: config comes back as an empty object and there's no safetensors field in the API response at all, which means no Hub-computed parameter count anywhere on the model page. That's normally where I'd paste in the Hub's number and move on. Here there isn't one, so I built it myself two independent ways — and it's the most interesting thing in this piece.

The parameter count, from bytes and from source

The README calls AuK "a 1.5B foundation model." The technical report's architecture section says the backbone has "approximately 1.5 billion" parameters. Neither says at what precision the released checkpoint is stored, and the Hub can't tell you either. So: auk_base.safetensors is 6,122,209,092 bytes. Divide by a dtype size and see which one the paper agrees with.

6,122,209,092 bytes ÷ 2 (bf16) = 3,061,104,546 params  ≈ 3.06B   — 2x too many
6,122,209,092 bytes ÷ 4 (fp32) = 1,530,552,273 params  ≈ 1.53B  — matches "~1.5B"

bf16 is the natural first guess — it's the model's own default inference dtype (--dtype bf16 in the CLI, torch_dtype=torch.bfloat16 for the frozen text encoder in infer_auk.py). It's off by almost exactly 2x. FP32 matches. And the code explains why: AukInfer.__init__ builds the model, then — in the release-day source — ran model = model.to(torch.float32) before loading the checkpoint weights, and only casts to bf16 at compute time via torch.autocast("cuda", dtype=self.dtype, ...) inside the sampling loop. The EMA checkpoint was exported at full precision; the bf16 you actually run it in is an autocast decision made at load time, not a property of the file on disk. Tencent shipped a 6.12 GB file for a model that only ever computes in bf16 — a bf16 export would have been 3.06 GB, half the download, for identical inference behavior.

I didn't want to stop at "the arithmetic is consistent with the paper's rounded number," so I counted the backbone from scratch, using nothing but the released config.yaml and the actual nn.Linear shapes in src/auk/model/modules.py:

# ckpts/AuK/config.yaml — model.arch, verbatim
arch:
  dim: 1536
  heads: 24
  ff_mult: 2
  text_hidden_dim: 2048
  num_layers: 10          # dual-stream (MMDiT) block count
  num_single_layers: 20   # single-stream (DiT) block count

With dim=1536, heads×dim_head = 24×64 = 1536 = dim, and ff_mult=2, every nn.Linear in DiTBlock and MMDiTBlock has a fixed, computable shape:

AdaLayerNorm:     dim → 6·dim         = 1536×9216 + 9216        = 14,164,992
Attention (self): to_qkv dim→3·dim, to_out dim→dim, 2 RMSNorms  =  9,443,456
SwiGLU FFN:       dim → 2·(dim·2), (dim·2) → dim, no bias       = 14,155,776
                                                    DiT block   = 37,764,224

MMDiTBlock duplicates every one of those — separate AdaLN, separate to_qkv/to_qkv_c, separate output projections, separate FFN per stream — and because context_dim is set equal to dim here, each duplicate is an exact copy of its DiT-block counterpart, not an approximation: 75,528,448 params, exactly 2× a DiT block. Ten MMDiT blocks and twenty DiT blocks then land at precisely the same total —

10 × 75,528,448  =  755,284,480   (dual-stream stack)
20 × 37,764,224  =  755,284,480   (single-stream stack)
                    ───────────
                  1,510,568,960   core transformer

— plus the small peripheral modules the two stacks share (time_embed, txt_proj/txt_norm, audio_embed's linear + two grouped 1D convs, the final AdaLayerNorm_Final, proj_out, and the ELMo-style layer-fusion weights), which add up to 19,969,669 more. Total: 1,530,538,629.

Correction, 18 September 2026: I should have read the header

Everything above was written without ever opening the checkpoint, and two of its numbers were wrong. You do not have to download 6.12 GB to check a safetensors file's contents — the first eight bytes are a little-endian u64 giving the length of a JSON header, and that header lists every tensor's name, dtype and shape. Two HTTP range requests, no weight bytes:

import json, struct, requests
URL = "https://huggingface.co/tencent/AuK/resolve/main/auk_base.safetensors"
n = struct.unpack("<Q", requests.get(URL, headers={"Range": "bytes=0-7"}).content)[0]
hdr = json.loads(requests.get(URL, headers={"Range": f"bytes=8-{8 + n - 1}"}).content)
print(n, len(hdr), {v["dtype"] for v in hdr.values()})          # 54568 420 {'F32'}
print(sum(__import__("math").prod(v["shape"]) for v in hdr.values()))  # 1530538629

All 420 tensors are F32. Not inferred from a byte count — declared, in the file. The fp32 conclusion was right for the right reason.

The exact count is 1,530,538,629, and it lands between the two estimates rather than on either:

1,530,552,273   file size ÷ 4               — 13,644 too high
1,530,551,393   my from-scratch derivation  — 12,764 too high
1,530,538,629   the header                  — measured

Both errors have identifiable causes, and the byte-division one is exact. 6,122,209,092 ÷ 4 divides the whole file, header included, by four. The header is 8 + 54,568 = 54,576 bytes, and 54,576 ÷ 4 = 13,644 — precisely the overcount. Subtract the header first and the shortcut becomes exact: (6,122,209,092 − 54,576) ÷ 4 = 1,530,538,629, to the parameter. So the original claim that the two methods agreed "to within 880 parameters" was true of each other and of neither of them, and the closeness was partly luck: two independent errors that happened to land within a thousand of each other.

What the header does vindicate is the structural derivation, which is the part that took the work. Every one of the ten dual-stream blocks is exactly 75,528,448 parameters and every one of the twenty single-stream blocks exactly 37,764,224 — the two numbers computed above from config.yaml and the nn.Linear shapes alone, reproduced to the parameter, and with them the 1,510,568,960-parameter core and the exact 2× ratio the 10-then-20 split is built around. The error was entirely in the peripheral modules: 19,969,669 in the file against my 19,982,433, a 12,764-parameter miss on a 20-million-parameter estimate. One piece of that is nameable — I described "33-parameter ELMo-style layer-fusion weights"; the file carries layer_weights of shape [36] plus a scalar layer_scale, so 37, and 36 is exactly the number of hidden layers in the Qwen2.5-Omni-3B Thinker whose stack it is weighting.

The backbone's parameter count is now settled at 1,530,538,629 rather than ≈1.531B, and so is AuK-Flash's: its checkpoint is byte-for-byte the same size with the same 420 tensors, the same shapes and the same all-F32 dtype, which is worth knowing — the distillation changed the weights, not the architecture or the export convention.

receiptscaptured 2026-09-18

The released checkpoint's own safetensors header settles what this article originally had to infer from the file's byte count. Every one of auk_base.safetensors' 420 tensors is F32 — the fp32 storage conclusion was right — and the header's shapes sum to 1,530,538,629 parameters. That is 13,644 below the file-size-divided-by-four estimate, because dividing the whole file by 4 also divides the 54,576-byte header by 4. AuK-Flash's checkpoint has the identical shape, tensor count and dtype.

artifacttensorsparametersdtypefile bytes
tencent/AuK — auk_base.safetensors4201,530,538,629F32 (all 420)6,122,209,092
└ 10 × MMDiT dual-stream block200755,284,480F323,021,137,920
└ 20 × DiT single-stream block200755,284,480F323,021,137,920
└ peripheral (embeds, norms, proj_out, fusion)2019,969,669F3279,878,676
tencent/AuK — vae.safetensors1,137159,299,797F32 (all 1137)637,322,604
tencent/AuK-Flash — auk_flash.safetensors4201,530,538,629F32 (all 420)6,122,209,092
Qwen2.5-Omni-3B Thinker, minus vision tower (resident)9244,034,780,160BF16 on disk8,069,560,320
└ thinker.model (36-layer LM body)4343,085,938,688BF166,171,877,376
└ thinker.audio_tower489637,676,544BF161,275,353,088
└ thinker.lm_head (allocated, unused)1311,164,928BF16622,329,856
└ thinker.visual (deleted at load)518668,684,288BF161,337,368,576

The block-level rows are the sums of the per-block tensor shapes, not a model of them: every one of the ten dual-stream blocks is exactly 75,528,448 parameters and every one of the twenty single-stream blocks exactly 37,764,224, which is what this article derived from config.yaml alone before the header was read. The `bytes` column for sub-rows is params × 4 (F32) or × 2 (BF16) and therefore excludes the header. Qwen2.5-Omni-3B's full repo is 5,537,120,672 parameters; AuK resides only the Thinker and deletes its vision tower, leaving 4,034,780,160.

method HTTP range reads against huggingface.co: bytes 0-7 give the safetensors header length as a little-endian u64, bytes 8..8+n-1 give the JSON header. Parameter counts are the product of each tensor's declared shape, summed; no weight bytes were downloaded. The Qwen2.5-Omni-3B figure is the same read over its three shards, restricted to the tensors AuK keeps at runtime (the Thinker, after infer_auk.py does `del thinker.visual`).
data /articles/auk-speech-editing/data/header-ledger.json (11 rows, 3.4 KB)
AuK backbone (config.yaml + src/auk/model/modules.py), 30 transformer blocksdim 1536 · 24×64 heads · ff×2
text streamaudio streamconcat -> single sequenceblock 1block 10/11 boundaryblock 30
block 10/30
dual-stream MMDiT block 10 — separate AdaLN, separate to_qkv/to_qkv_c, separate output projection, separate SwiGLU FFN per stream. One joint attention call: Q/K/V from both streams are concatenated, then attended together. ~75.5M params/block.
depth-parameters so far755.3M / 1.51B

Because AuK sets context_dim equal to dim, every module MMDiT duplicates is an exact copy of its DiT counterpart — so a dual-stream block costs precisely 2× a single-stream block, computed from the same config.yaml shapes shown above: 75.5M against 37.8M. Ten of the first against twenty of the second lands both stacks at exactly 755.3M each — the 10-then-20 layer split isn’t a round number picked for its own sake, it’s a 50/50 compute split between fusing the two modalities and refining the merged one.

What "1.5B foundation model" doesn't include: the frozen Qwen2.5-Omni-3B semantic encoder — a separate, required download (hf download Qwen/Qwen2.5-Omni-3B), not bundled in tencent/AuK at all — and AuK-Flash, which lives in its own repo (tencent/AuK-Flash) with its own checkpoint. vae.safetensors (637,322,604 bytes, 159,299,797 params across 1,137 all-F32 tensors, by the same header read) is bundled and does count toward the 6.76 GB usedStorage, but it's a separate module the README explicitly calls out: "The model checkpoint contains the diffusion transformer and layer-fusion weights. The MLLM encoder and VAE are loaded from separate files at runtime, so missing text_encoder.* keys during checkpoint loading are expected." Running AuK end to end means loading 1,689,838,426 parameters from the tencent/AuK repo plus an encoder you have to fetch yourself — a materially bigger footprint than the headline number suggests, none of it wrong, all of it worth knowing before you plan a deployment around "1.5B."

And "~3B encoder," which is what I originally wrote here, was wrong. Qwen/Qwen2.5-Omni-3B is 5,537,120,672 parameters in total, and what infer_auk.py keeps resident is not the 3B in the name. It instantiates Qwen2_5OmniThinkerForConditionalGeneration and then deletes exactly one thing:

# src/auk/infer/infer_auk.py — AukInfer.__init__
thinker = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
    text_encoder_config.text_encoder_path,
    torch_dtype=torch.bfloat16,
)
# keep the full multimodal Thinker (text + ref_audio); drop the unused vision tower
if thinker.visual is not None:
    del thinker.visual
    thinker.visual = None

Summing the shapes in Qwen2.5-Omni-3B's three shard headers and removing only thinker.visual (668,684,288 parameters) leaves 4,034,780,160 resident: the 3,085,938,688-parameter 36-layer LM body that the "3B" refers to, plus a 637,676,544-parameter audio_tower AuK genuinely needs for reference-audio conditioning, plus a 311,164,928-parameter lm_head it allocates and never uses, since it reads hidden states rather than logits. Total resident across all three stacks: 5,724,618,586 parameters to run a model marketed at 1.5B. That number is about to matter a great deal.

What 1.95 million hours actually means

The abstract's data claim: "approximately 3.03 billion instruction–audio instances and 1.95 million hours of effective supervision" across five task families. 1.95M hours is 222 years of audio — the kind of number that's supposed to make you stop reading and be impressed, so it's worth doing the denominator check the number deserves.

3.03 billion instances over 1.95 million hours averages to about 2.3 seconds of audio per instance (1.95M × 3600 ÷ 3.03B). That's a plausible average for short TTS/editing utterances, so the two headline numbers aren't obviously inconsistent with each other. What I couldn't find anywhere in the paper is a definition of "effective supervision hours" itself: does an hour of real speech that gets reused across, say, a zero-shot-TTS instance, a pitch-edit instance, and an enhancement instance count as one hour, or three? The paper states the two headline numbers and moves on. I looked for the defining sentence and didn't find one — which matters, because "1.95 million hours of unique audio" and "1.95 million hours of task-instance audio, summed across every augmentation of the same underlying clips" are very different claims about how much real speech went into this model, and only one of the two supports "222 years of audio" as a statement about data collection rather than data multiplication.

The task-family split (paper's Table 1, pre-training stage 2 sampling probabilities) does check out arithmetically — it sums to exactly 100%:

Task familyShare
Speech Generation28.10%
Enhancement & Separation23.17%
Content Editing23.02%
Paralinguistic Editing21.75%
Acoustic Editing3.96%

The paper is candid, if qualitative, about how much of this is synthetic: emotion editing is trained on Qwen3-TTS-CustomVoice-synthesized pairs, timbre editing on the X-VC corpus, nonverbal editing on F5-TTS-generated audio, content editing on masked-infilling targets, and enhancement/separation on "independently sampled degradations" applied to clean speech. That's most of the five families built at least partly from another model's output rather than from a microphone. No numeric real-vs-synthetic split is given anywhere I could find — another place where "if you could not verify something, write that" applies directly.

AuK-Flash: 4 steps, no CFG, and a 4.5x that should have been bigger

AuK-Flash is a distilled release for fast inference: 4 fixed sampling steps, classifier-free guidance switched off (CFG=0), versus the full model's default of 32 steps with CFG=2.0 — both defaults are the actual code's, not paraphrased:

# src/auk/model/cfm_edit.py — CFMEdit.sample()
def sample(
    self,
    cond, text, duration, *,
    steps=32,
    cfg_strength=1.0,
    ...
# src/auk/infer/infer_cli.py — CLI defaults
--nfe    32     # number of function evaluations (ODE steps)
--cfg    2.0     # classifier-free guidance strength
# README.md
AuK-Flash use 4 fixed time steps and set CFG=0.

CFG at inference isn't a second forward pass bolted on afterward — flux2_edit.py's cfg_infer path concatenates the conditional and unconditional inputs into one batch (x = torch.cat((x_cond, x_uncond), dim=0)) and runs one forward over twice the batch, which costs about 2× a non-CFG step. So the naive accounting for "how much cheaper is Flash" is steps × CFG-multiplier on both sides: 32×2 = 64 units for the full model, 4×1 = 4 units for Flash — a 16× step-cost ratio. The paper's abstract states a 4.5× wall-clock speedup "under matched conditions," without saying what those conditions are (no batch size, no hardware, no sequence length disclosed for that specific number). 4.5× is real work and a genuine win. It's also a third of what the step math alone predicts, and the paper doesn't say where the rest goes.

cfm_edit.py sample() · README (32 NFE + CFG=2.0 vs Flash’s 4 steps, CFG=0)reconstructed, not paper-stated
At 4 steps with CFG off: naive step-cost ratio 16.0x, but with the fixed text-encoder-plus-VAE-decode cost included, estimated wall-clock speedup 4.5x.AuK (full)77.1 unitssteps=4, CFG off17.1 units fixed (text encoder + VAE decode, ~13.1 units)    diffusion steps (scales with NFE × CFG)
steps=4
naive step-cost ratio
16.0×
with fixed cost included
4.5×

Set steps to 4 and CFG off — AuK-Flash’s actual release config — and the reconstructed speedup lands at 4.5×, matching the paper’s stated 4.5×. The naive ratio for the same setting is 16.0×, because dropping CFG halves the cost of each step and cutting 32 steps to 4 divides by 8 more. The gap between the two numbers only closes if something costs the same no matter how many diffusion steps you take — and the code shows exactly one thing that fits: encode_text() runs the frozen Qwen2.5-Omni-3B thinker once per request, before the sampling loop, and its output is cached across every step. Drag steps back toward 32 with CFG still off and both numbers fall toward 1× — there is less step-count left to cut — but the realized speedup stays below the naive one at every setting, because the 13.1-unit fixed cost never goes away; it’s just a shrinking share of a shrinking total.

My reconstruction, not the paper's: something has to cost the same regardless of step count or CFG for a 16× ratio to land at 4.5×, and cfm_edit.py shows exactly one thing built that way. encode_text() runs the frozen Qwen2.5-Omni-3B thinker exactly once per request — before the ODE loop starts — and its output is threaded through every subsequent step via the transformer's cache=True argument, never recomputed. VAE decode is likewise one pass at the end. Solving "what fixed cost, added to 64 units and to 4 units, turns a 16× ratio into a measured 4.5×" gives roughly 13.1 units of fixed cost relative to one non-CFG backbone step — plausible for a full forward pass through a 3B-parameter encoder plus a VAE decode, both amortized over just 4 diffusion steps instead of 64. That's an estimate built to match the disclosed 4.5×, not an independently measured number; I'm flagging it as reconstruction because the paper gives no latency breakdown to check it against directly.

Leading on generation and editing, and where "competitive" is doing real work

The abstract's own phrasing: "leading performance on zero-shot and instruction-controlled speech generation... [and] general speech editing... while remaining competitive on signal-level restoration tasks." That's three different verbs for three different task groups in one sentence, and the paper's own figures back the distinction up with real numbers rather than just softening the tone for the harder category.

Generation and editing — AuK leads outright, against strong, named baselines:

BenchmarkMetricAuKAuK-FlashBest specialist baseline
Seed-TTS-EvalWER ↓2.65%2.85%Qwen3-TTS 3.07%
Seed-TTS-EvalSIM ↑0.7950.790Seed-TTS 0.778
InstructTTSEval (DSD, ZH)acc. ↑83.37%78.80%Qwen3-TTS-VD 81.10%
MMAE-SpeechEMR ↑13.8512.44Ming-UniAudio 7.04
SpeechEditBenchscore ↑49.7346.66Ming-UniAudio 28.70
Ming-Freeform-Audio-Editscore ↑88.5487.55Ming-UniAudio 76.99

Signal-level restoration is where "competitive" earns its precision. On DNSMOS-OVRL — the automatic perceptual-quality metric an enhancement specialist is built to optimize — AuK loses to a dedicated enhancement model on both datasets where the comparison is made directly:

BenchmarkMetricAuKRE-USE (specialist)Who wins
DNS ChallengeDNSMOS-OVRL ↑3.353.38RE-USE
DNS ChallengeUTMOS ↑3.863.69AuK
CHiME-4DNSMOS-OVRL ↑3.283.33RE-USE
CHiME-4WER ↓7.98%10.71%AuK
Libri2Mix (separation)DNSMOS-OVRL ↑3.283.24 (MossFormer2-SS)AuK
Libri2Mix (separation)WER ↓9.12%9.34% (MossFormer2-SS)AuK

That's the honest shape of "competitive": on the two enhancement benchmarks, AuK loses the automatic-MOS metric to the specialist by a small but real margin (0.03–0.05 DNSMOS-OVRL) while winning decisively on intelligibility (WER) and a second naturalness metric (UTMOS) — a real, disclosed trade rather than a clean win dressed down for modesty. On the one separation benchmark shown, AuK actually edges the specialist on every visible metric, so "competitive" undersells that particular comparison if anything. Both figures come straight off the paper's own bar chart (assets/performance.png), reproduced below.

Grouped bar charts across three panels. Panel a, Speech Generation Capability: AuK and AuK-Flash lead Qwen3-TTS, Seed-TTS, VoxCPM2, MOSS-VoiceGenerator, Mimo-Audio and Qwen3-TTS-VD on Seed-TTS-Eval WER and SIM and on InstructTTSEval DSD. Panel b, Speech Editing Capability: AuK and AuK-Flash lead Ming-UniAudio and Step-Audio-EditX on MMAE-Speech EMR, SpeechEditBench, and Ming-Freeform-Audio-Edit by a wide margin. Panel c, Speech Enhancement and Separation Capability: paired DNSMOS-OVRL and UTMOS bars for AuK, AuK-Flash, SAM-Audio-Large, RE-USE, and MossFormer2-SS across DNS Challenge, CHiME-4, and Libri2Mix, where AuK's DNSMOS-OVRL bar sits slightly below the specialist RE-USE on the first two benchmarks and slightly above MossFormer2-SS on the third.
Every headline benchmark panel from the model card, reproduced at full resolution — including the two DNSMOS-OVRL bars where AuK sits below its specialist baseline (paper, performance figure).

Architecture: an MLLM for meaning, a shared VAE, an SD3-shaped backbone

Semantic conditioning. AuK doesn't use Qwen2.5-Omni-3B's text output — it uses its hidden states, aggregated ELMo-style: every one of the encoder's internal layers gets a learned softmax weight and a per-layer LayerNorm, and the weighted sum (times one more learned scalar) becomes the conditioning signal. That's a real, small piece of code, not a paraphrase of the paper's Equation 2:

# src/auk/model/cfm_edit.py — CFMEdit.encode_text()
stacked = torch.stack([F.layer_norm(h, [d_llm]) for h in all_hidden_states[1:]], dim=0)
weights = F.softmax(self.layer_weights, dim=0)
hidden = (stacked * weights[:, None, None, None]).sum(dim=0) * self.layer_scale

The rationale for using an LLM at all rather than a text-only encoder: the same interface has to carry a spoken-language instruction ("replace X with Y"), a free-text voice description, and reference audio, and a multimodal model already knows how to attend across all three without three separate encoders.

The VAE is trained jointly on speech, general audio, and music at a stated 6:3:1 sampling ratio — one shared latent space instead of a speech-only codec, which is what lets one model do TTS, music source separation, and target-speaker extraction without swapping representations. It maps 24 kHz waveform to 64-dimensional latents at 50 Hz (a 480× time-domain downsample: downsample_rate: 480 in config.yaml, confirmed by the encoder's downsample_rates: [2,2,2,3,4,5] product), and decodes through a BigVGAN-style stack of six transposed-convolution blocks.

The backbone's dual-stream-then-single-stream shape is the SD3/FLUX MMDiT pattern, and it shows up here for the same reason it shows up in image diffusion: text and audio start with different statistics and need their own normalization and feed-forward transforms early on, but attention is what actually lets one modality condition the other, so the joint-attention call is shared from block one — only the surrounding per-stream layers (AdaLN, projections, FFN) stay separate. Once ten blocks of that have let the two streams settle into a shared representation, there's no more reason to keep them apart, so the remaining twenty blocks run as one sequence through cheaper, shared weights. The interactive above walks through why that 10-then-20 split isn't arbitrary: a dual-stream block costs exactly 2× a single-stream one here (because context_dim == dim), so 10 of the expensive kind and 20 of the cheap kind spend exactly the same total parameters on each half — a genuinely balanced design, not a round number.

Three-panel architecture diagram. Panel a: the overall AuK pipeline, showing text semantic condition tokens, optional audio semantic condition tokens, optional reference audio condition tokens, and noisy VAE condition tokens feeding through a stack of Dual Stream MMDiT blocks repeated M times, then Single Stream DiT blocks repeated N times, into a frozen VAE decoder that outputs generated audio; a frozen large language model with per-layer hidden states h_i feeds a learned layer-wise weight sum and learned scale into the text condition tokens. Panel b: the internal structure of a Dual Stream MMDiT block, with separate layer norm, scale-and-shift, and SwiGLU FFN branches for a semantic stream x and an acoustic stream y, joined by one joint attention module over concatenated Q, K, V. Panel c: the internal structure of a Single Stream DiT block, with one layer norm, self attention, and SwiGLU FFN operating on the concatenated semantic and acoustic streams.
AuK's full architecture: the frozen LLM and VAE stages, the dual-stream MMDiT block's separate-but-jointly-attending design, and the single-stream DiT block that follows it (paper, architecture figure).

Two post-training methods for two different problems

AuK's post-training splits by task family, and the split is deliberate rather than incidental: editing gets human-feedback preference optimization; generation gets automatic-reward reinforcement learning. The paper's own reasoning is that these two task families fail differently. Zero-shot TTS has a cheap, reliable way to score a candidate automatically — did the ASR transcript match, does the speaker embedding match the reference — so a reward model can be built and RL can run at scale. Whether an edit "sounds natural" or "changed the right amount" has no such automatic proxy, and the paper states plainly that neither a broad-enough preference dataset nor a ready-made reward model existed for it, so they collected human ratings directly instead.

Editing: Diffusion-DPO with a listwise (LiPO) objective. For each editing instruction, annotators rate 10–20 candidate outputs on a 3-level scale (0 = failure, 1 = partial, 2 = success); groups where every candidate got the same rating are discarded as carrying no signal. What's left: 818 informative groups, 9,080 rated candidates. The loss compares candidates pairwise within {1≻0, 2≻0, 2≻1} relations, weights the strongest relation (2≻0) at 1.0 and the other two at 0.5, and trains for 104 optimizer updates from the pre-training checkpoint.

Generation: Flow-GRPO with two different automatic reward models, one per generation sub-task. Zero-shot TTS's reward blends a tolerant word-error-rate (homophone substitutions forgiven — a Chinese homophone swap isn't a pronunciation failure) with a strict WER, plus cosine similarity between speaker embeddings, standardized only among candidates that already got the content right (so a good voice match can't paper over wrong words). Instruction-controlled TTS's reward is a dedicated style-consistency judge — a Qwen2.5-Omni-7B model trained on 30,000 balanced positive/negative examples, queried by majority vote. Both run for 500 optimizer updates, 16 candidates sampled per prompt, with rollout sampling stochastic only in a 6-step low-SNR window and deterministic everywhere else, to keep variance and cost down.

Update — 18 September 2026

Nine days on, the release has moved in three ways, and the most interesting of them is a direct consequence of the fp32 finding above.

The repo's own News block is the timeline:

DateWhat landed
2026/09/09Open-sourced: code, weights, Gradio demo Space, ModelScope studio. SGLang-Omni day-0 support.
2026/09/13CPU offload for CUDA inference, and MLX on Apple Silicon on the feat/mlx-apple-silicon branch.
2026/09/13AuK named the end-to-end baseline for the Single Model Track of the ICASSP 2027 Audio Editing Challenge.
2026/09/16Encoder memory cut by ~7.5 GiB (PR #19).

The 7.5 GiB was the fp32 cast, spent twice

PR #19, authored by Zhikang Niu and merged 16 September, changes exactly one line of src/auk/infer/infer_auk.py:

-model = model.to(torch.float32)
+model.transformer.to(torch.float32)

That is the whole diff. Its title says what it is: "Cast only the DiT backbone to fp32 to avoid the Qwen encoder upcast."

Read it against the section above and the mechanism is obvious in hindsight. CFMEdit is a container: it holds the Flux2Edit backbone and the frozen Qwen Thinker and the VAE. The Thinker had just been loaded with torch_dtype=torch.bfloat16 — deliberately, because it is frozen and never trained — and then the module-wide .to(torch.float32), aimed at the backbone, swept it back up to fp32 and doubled it. The cast that let me prove the checkpoint was fp32 was the same cast that was costing 7.5 GiB of VRAM on every single run.

The numbers in Tencent's PR description are reproducible from the header reads above, which is the best kind of check:

safetensors headers · infer_auk.py before and after PR #19measured, not modelled
Resident weight bytes before and after PR #19. Before: encoder 15.031 GiB plus DiT 5.702 GiB plus VAE 0.593 GiB. After: encoder 7.515 GiB, the other two unchanged. Tencent measured loading peaks of 21.44 and 13.97 GiB respectively.before PR #19model.to(torch.float32) — encoder upcast to fp3221.33 GiBvs 21.44 measuredafter PR #19model.transformer.to(torch.float32) — encoder stays bf1613.81 GiBvs 13.97 measured Qwen2.5-Omni Thinker, vision tower deleted    Flux2Edit DiT (fp32 either way)    BigVGAN-Flow VAE
encoder, fp32
15.031 GiB
encoder, bf16
7.515 GiB
difference
7.515 GiB
Tencent’s figure
7.516 GiB

Three parameter counts and two dtypes reproduce Tencent’s own before/after loading peaks to within 0.16 GiB — 21.33 GiB against a measured 21.44, and 13.81 GiB against a measured 13.97. The gap in each case is allocator slack, not a missing module. Note what does not move: the DiT stays at 5.702 GiB in both rows, because fp32 really is the checkpoint’s storage dtype and the fix deliberately keeps it. The entire 7.52 GiB came out of a frozen encoder that had been loaded in bf16 and then cast back up to fp32 by a blanket module-wide .to() that was only ever meant for the backbone.

Their measured encoder storage, before and after, is 15.031 GiB and 7.515 GiB. The 4,034,780,160 resident Thinker parameters I counted come to 15.0307 GiB at four bytes each and 7.5153 GiB at two. Their measured model-loading peaks, 21.44 GiB and 13.97 GiB on an H20, are within 0.16 GiB of the sum of the three checkpoints at their respective dtypes. Nothing here is a heuristic: the saving is exactly one module, counted once instead of twice.

CPU offload, and a README table that no longer describes the code

The 13 September release added --cpu_offload (also cpu_offload=True on AukInfer, and a checkbox in the ComfyUI loader). It uses accelerate's cpu_offload_with_hook to page the Thinker and the DiT in and out one at a time while deliberately keeping the VAE GPU-resident, and the README publishes this table, torch.cuda.max_memory_allocated on one A800-SXM4-80GB with bf16 inference:

ModelInputoffload offoffload onsaved
AuKtext only, 1.5 s out24.78 GiB16.75 GiB8.03 GiB (32.4%)
AuK5 s reference audio25.00 GiB16.98 GiB8.02 GiB (32.1%)
AuK-Flashtext only, 1.5 s out24.77 GiB16.75 GiB8.02 GiB (32.4%)
AuK-Flash5 s reference audio24.97 GiB16.98 GiB7.99 GiB (32.0%)

Those numbers are now stale, and the README does not say so. Every row was measured on 12 September, two commits before the encoder fix. PR #19's own table puts the post-fix inference peak at ~18 GiB without offload and ~10.5 GiB with it. The commit that announced the fix, 871bf3d, is one line long: it adds the News bullet and re-measures nothing. So the README currently tells you in its News section that the encoder got 7.5 GiB cheaper and then, four hundred lines down, hands you a memory table from before that happened. If you are sizing a box off that table you are provisioning about 7 GiB you no longer need.

Worth being precise about what CPU offload buys and why it is not the same 7.5 GiB. Offload does not shrink anything; it stops two stacks from being resident simultaneously. Peak becomes roughly the larger of them plus the VAE, which pre-fix was the 15.03 GiB encoder — which is why the pre-fix saving, 8.03 GiB, is close to the 5.70 GiB DiT plus its activations rather than to anything about precision. Post-fix, with the encoder at 7.515 GiB and the DiT at 5.702, the two stacks are nearly the same size and offload has much less to trade. The two fixes are complementary, and they compose: the second one shrinks what the first one was paging.

MLX on Apple Silicon — a branch, not a merge

The MLX port is real and unusually well documented, but note where it lives: feat/mlx-apple-silicon, not main. It is a separate backend under src/auk_mlx/ that imports nothing from src/auk/, behind an optional .[mlx] extra, with a one-time conversion step. Its published numbers, on an M4 Pro with 48 GB, base model, 32 steps, 6 s of audio:

ConfigurationPeak memoryRTF
fp3224.3 GB6.05
fp32, --sequential15.4 GB6.05
8-bit9.1 GB6.02
8-bit, --sequential6.1 GB6.02

Two things in that table are worth reading carefully. First, --sequential is the MLX analogue of --cpu_offload and it is free here — under unified memory, evicting a stack is a deallocation rather than a device-to-host copy, which is why it costs nothing in RTF while the CUDA path's offload does pay for transfers. Second, and stated plainly on the branch: "Quantization buys memory, not speed" — wall clock is flat to within 0.5% across fp32, 8-bit and both sequential variants, because a 32-step ODE over a 1.5B backbone is compute-bound, not bandwidth-bound. That is the honest version of a claim most quantization announcements round the other way. An RTF above 6 also means this is not a real-time path on a laptop; it is a "works at all without CUDA" path.

The branch's own arithmetic checks out where I could test it. "About 28 GB of fp32 MLX weights" after conversion: 0.637 (VAE) + 6.122 (base DiT) + 6.122 (Flash DiT) + 14.894 (Thinker) = 27.78 GB. That Thinker figure is the tell — 14.894 GB is 3,723,615,232 parameters at four bytes, which is the 4,034,780,160-parameter resident Thinker minus its 311,164,928-parameter lm_head. The MLX converter drops the output projection that the PyTorch path allocates and never uses. One number I could not reconcile as written: the branch describes the VAE as "only 147M parameters," where the released vae.safetensors header sums to 159,299,797. The gap is exactly the 12,418,176-parameter flow.* stack — 159,299,797 − 12,418,176 = 146,881,621 — so "147M" describes the converted VAE after the normalizing-flow module is dropped, not the file you download. Correct for what it measures; it just is not measuring the artifact.

The demo Space is running the old code

One last thing, because it is checkable in thirty seconds and nobody seems to have. The official demo Space vendors its own copy of the inference source, and that copy has not been touched since 10 September — before the offload PR and before the encoder fix. Line 103 of spaces/tencent/AuK/src/auk/infer/infer_auk.py still reads model = model.to(torch.float32); the string cpu_offload appears in it zero times; and app.py builds two AukInfer engines at startup, base and Flash, each with its own full-precision copy of the Thinker, then asks ZeroGPU for a large slice (spaces.GPU(duration=120, size="large")). The demo everyone clicks to evaluate this model is therefore running the most expensive configuration the project ships, two days after the project stopped shipping it. It works — the Space is up — but "AuK needs 25 GB" and "the AuK demo needs a big GPU" are both statements about a loader that main has already replaced.

What checks out, plainly

Not every number here needed a caveat. The task-family percentages sum to exactly 100%. The zero-shot and instruction-TTS benchmark numbers in the bar chart match the paper's own tables to two decimal places, and where the chart shows an averaged ZH/EN score I couldn't find broken out anywhere else, recomputing the average from the paper's per-language numbers reproduces the chart's bar to the hundredth (82.49 from 83.37/81.60, 81.75 from 81.10/82.40, and so on — five-for-five across the InstructTTSEval panel). The default sampling config in the actual inference code (--nfe 32 --cfg 2.0) matches what the paper describes as the full model's setting, so the benchmark numbers weren't measured on a different configuration than what ships. And the backbone's parameter count — the thing with no Hub metadata to check against — turned out to be the most verifiable number in the whole report, once you're willing to read the checkpoint itself: 1,530,538,629 parameters across 420 tensors, every one of them F32.

What doesn't check out cleanly: no stated definition for "effective supervision hours," no numeric real-vs-synthetic data split, no disclosed hardware or batch size for the one benchmark (the 4.5× speedup) where "matched conditions" is the entire methodology section, and a released repo whose headline parameter count describes 27% of what you actually need in memory to generate anything — 1.53B of 5.72B, once the VAE and the resident Thinker are counted honestly.

And what this article got wrong, since the standard has to cut both ways: the byte-division parameter estimate and the from-scratch derivation agreed with each other and not with the file, because I compared two estimates instead of opening the header; the layer-fusion weight vector is 36 wide, not 33; and the semantic encoder is 4.03B parameters resident, not "~3B." All three are corrected above, in place, with the original reasoning left standing so the failure mode is visible. The structural derivation — that every dual-stream block is exactly twice a single-stream block, and that ten of one and twenty of the other spend identical parameter budgets — is reproduced by the header to the parameter, which is the part that was worth doing.

Cite this article

For attribution, please use the following reference or BibTeX:

Satyajit Ghana, "AuK: 6.12 GB, an empty Hub config, and a parameter count that only balances in FP32", ai.thesatyajit.com, September 2026.

bibtex
@misc{ghana2026aukspeechediting,
  author = {Satyajit Ghana},
  title  = {AuK: 6.12 GB, an empty Hub config, and a parameter count that only balances in FP32},
  url    = {https://ai.thesatyajit.com/articles/auk-speech-editing},
  year   = {2026}
}
share