cortexelus commited on
Commit
91ec853
Β·
verified Β·
1 Parent(s): 46e266c

small DiTs (sm-music, sm-sfx) tflite: ship the w8a8 rung ladders

Browse files

Replace the varlen small DiTs with from-repo rung ladders: dit_{fp32,w8a8}.tflite (8 rungs to the 2-min max 1292, gce-outside, gcond bundled). Flat ~0.8 GB RAM vs the old varlen's ~3.8 GB at 2 min (4.5x less), ~equal quality to w8a8-dyn (cos>=0.996), cache-safe. Default is now w8a8 for every family. Old tiers (fp32/w16a32/w8a32/w8a8-dyn) -> legacy/. Pairs with stable-audio-3 PR #118. (No .xnnwc: the small w8a8 packs in ~1s.)

tflite/README.md CHANGED
@@ -63,14 +63,16 @@ prompt ─▢ T5Gemma encoder ─▢ DiT pingpong sampler ─▢ SAME-S/L decode
63
 
64
  Precision is split between the **DiT** and the **SAME codec** β€” they react to quantization oppositely:
65
 
66
- - **DiT** (`--dit-precision`, or the global `--precision`):
67
- - **medium**: **`w8a8` (default)** Β· `fp32` β€” a from-repo **rung ladder** (12 fixed lengths; see *Medium DiT β€” the
68
- N=12 rung ladder* below). `w8a8` is **cache-safe int8**: ~equal quality to the old `w8a8-dyn` (cos 0.996 per
69
- forward) but with **flat RAM (~2–4 GB vs the old 3–40 GB)**, faster at the rung sizes, and a single file.
70
- int8 on the DiT is still a *different* (not worse) sample β€” the 8-step sampler is chaotically sensitive β€” so
71
- `w8a8` is the shipped default and `fp32` is the bit-exact reference (and the pick for LoRA merges or an A/B).
72
- - **sm-music / sm-sfx**: `fp32` Β· `w16a32` Β· `w8a32` Β· `w8a8-dyn` β€” still the dynamic-shape varlen graphs
73
- (wXaY = weight/activation bits; "16" = fp16), until they are rung-built too.
 
 
74
  - **Codec** (`--decoder-precision` / `--encoder-precision`): `fp32` Β· **`w8a8` (default)**. The SAME enc/dec ship as
75
  static **rung** models β€” see [docs/RUNGS.md](docs/RUNGS.md). They run once on a fixed latent, so int8 is
76
  quality-free: a real-music round-trip is identical to fp32. `w8a8` is ~3Γ— faster and ~half the RAM on an
@@ -78,33 +80,30 @@ Precision is split between the **DiT** and the **SAME codec** β€” they react to
78
  codec variants added no audible quality and are retired to `tflite/same-*/legacy/`.)
79
 
80
  `--precision` is a global default that sets the DiT directly and maps to the codec (`fp32`β†’`fp32`, any int8β†’`w8a8`).
81
- When omitted, the DiT precision is **per-family β€” medium `w8a8`, the small DiTs `fp32`** (they have no rung `w8a8`
82
- yet) β€” and the **codec is `w8a8`**. All variants keep the full feature set (any `--seconds`,
83
- odd/even L, batched/sequential CFG) and lazy-download from HuggingFace on first use.
84
 
85
  | component | precisions | size | notes |
86
  |---|---|---|---|
87
- | **DiT β€” medium** (rung) | **`w8a8`** (default) / `fp32` | 1.6 / 5.8 GB | N=12 rung ladder; `w8a8` = cache-safe int8, **flat ~2–4 GB RAM**, β‰ˆ old `w8a8-dyn` quality; old `w16a32/w8a32/w8a8-dyn` β†’ `legacy/`; needs litert β‰₯ 2.2.0 |
88
- | **DiT β€” small** | `fp32` / `w16a32` / `w8a32` / `w8a8-dyn` | sm 1.8 / 0.9 / 0.45 / 0.45 GB | varlen (not yet rung); int8 β†’ a *different* sample |
89
  | **codec** (rung) | `fp32` / `w8a8` | SAME-S enc/dec 0.22 / 0.25 β†’ w8a8 0.06 / 0.09 GB Β· SAME-L 1.7 / 1.8 β†’ w8a8 0.47 / 0.52 GB | w8a8 quality-free, ~3Γ— faster + half RAM; needs litert β‰₯ 2.2.0 |
90
  | T5Gemma | fp16 | 0.6 GB | single-precision |
91
 
92
- *(int8 weights are per-channel; `w8a8-dyn` also quantizes activations, its int8Γ—int8 matmuls being what make it
93
- faster-than-fp32 on the DiT. Codec rungs use dynamic int8 too, but there the extra activation-quant error is far
94
- below the AE's own loss. DiT speed factors measured on an Apple M4 Pro, XNNPACK, 8 threads.)*
95
 
96
- One CFG note (small DiTs, varlen): batched vs sequential CFG is bit-identical at `fp32`/`w8a32`, ~80 dB apart at
97
- `w16a32`, but under `w8a8-dyn` the batch=2 invoke shares activation scales across the cond/uncond rows β†’ a
98
- *different plausible sample*. Pass `--no-cfg-batched` with `w8a8-dyn` when you need run-to-run reproducibility.
99
- (The medium rung DiT runs batch=1 / sequential CFG, so this doesn't apply to it.)
100
 
101
- ### Medium DiT β€” the N=12 rung ladder (replaces the old varlen tiers)
102
 
103
- The medium DiT now ships as a **static rung ladder**, built from the repo's own PyTorch model (see [build/](build/README.md)),
104
- not a dynamic-shape varlen graph. A render of length *L* runs on the **smallest of 12 fixed rungs β‰₯ L** in one forward,
105
- with the extra positions masked out of self-attention β€” so the result is **exact, not tiled**. Both `fp32` and the
106
- cache-safe int8 `w8a8` are single self-contained files (`dit_{fp32,w8a8}.tflite`; the per-step global-cond preamble
107
- rides inside as a `gcond` signature, so there is no extra file to ship).
108
 
109
  **Why it matters β€” RAM.** The old varlen graph materialized the full `[heads, S, S]` attention, so peak RAM grew with
110
  length and hit **~40 GB at the 380 s maximum** β€” unusable on a normal machine. The rung ladder's static shapes let
@@ -114,14 +113,19 @@ rung sizes and β‰ˆ equal quality to the old `w8a8-dyn` (cos 0.996 per forward).
114
 
115
  ![Medium DiT peak RAM β€” rung ladder vs varlen](docs/img/dit-ram-rung-vs-varlen.png)
116
 
117
- **The 12 rungs** (in latents; audio β‰ˆ L / 10.8 s, 4096 = 380 s max):
118
- `8, 48, 96, 192, 416, 704, 1056, 1496, 2040, 2824, 3536, 4096`.
119
- The short rungs β€” **8 (β‰ˆ0.7 s) Β· 48 (β‰ˆ4.5 s) Β· 96 (β‰ˆ8.9 s) Β· 192 (β‰ˆ18 s)** β€” keep tiny clips fast (a 1 s clip no longer
120
- pays for an 18 s render). A length *between* rungs pads up to the next one. The ladder is configurable via `DIT_LADDER`
121
- in [`build/build_dit.sh`](build/build_dit.sh).
122
 
123
- The old medium tiers (`w16a32`, `w8a32`, `w8a8-dyn`, and the previous varlen `fp32`) are retired to
124
- **`tflite/sa3-m/legacy/`** on HuggingFace.
 
 
 
 
 
 
 
 
125
 
126
  ```bash
127
  # medium DiT β€” defaults: w8a8 rung DiT + w8a8 SAME-S codec (the `--decoder` now defaults to same-s)
@@ -146,9 +150,9 @@ codec's precision maps directly to audio fidelity; the DiT's changes *which* sam
146
  you get; the encoder's affects a2a/inpaint latents):
147
 
148
  ```bash
149
- # fastest DiT, reference-quality codec
150
  ./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s \
151
- --precision fp32 --dit-precision w8a8-dyn
152
  ```
153
 
154
  ## Install
@@ -264,9 +268,8 @@ radio with pre-generation, loop, hotswap, MP3/WAV save).
264
  first run it offers to install the UI-only extras (`gradio`, `pillow`,
265
  `soundfile`) into `.venv`. In the browser:
266
 
267
- - **Precision** dropdown next to the model picker β€” `fp32` (default) / `w16a32`
268
- / `w8a32` / `w8a8-dyn`, applied to the DiT; the codec maps it to its rung tier
269
- (`fp32` β†’ fp32, any int8 β†’ w8a8).
270
  - **LoRA** panel β€” upload a `.safetensors` adapter or pick one from
271
  `loras/<model>/` (e.g. `loras/sa3-medium/`), set its strength, stack several;
272
  settings are remembered per model. The panel is disabled under a quantized
@@ -310,8 +313,8 @@ For sub-realtime latency on a supported device, prefer the GPU siblings:
310
  | `--seed` | random | Set for reproducibility; the chosen seed is printed at the end |
311
  | `--cfg` | 1.0 | Guidance scale; 1.0 = off, >1 toward prompt, <1 toward uncond. β‰ 1 runs cond+uncond each step |
312
  | `--apg` | 1.0 | Adaptive Projected Guidance; only matters when `--cfg β‰  1` |
313
- | `--cfg-batched` | on | When `--cfg β‰  1`, run cond+uncond as one batch=2 invoke on the variable-batch DiT (~7–29% faster on Apple-Silicon AMX). `--no-cfg-batched` β†’ sequential batch=1 dual-pass. Bit-identical at `fp32`/`w8a32` β€” see [Precision variants](#precision-variants---precision) for `w16a32`/`w8a8-dyn` |
314
- | `--lora` | β€” | A `.safetensors` LoRA adapter (SA3-native/underfit or PEFT) merged into the DiT, optional `strength=S`; repeat to stack. Requires `--dit-precision fp32` or `w16a32`. See [LoRA](#lora) |
315
  | `--lora-strength` | 1.0 | Default strength for `--lora` adapters without their own `strength=` |
316
  | `--init-audio` | β€” | WAV (any format via ffmpeg) input for audio-to-audio / inpaint |
317
  | `--init-noise-level` | 1.0 | Οƒmax; 0.4–0.8 typical for variation, 1.0 = full regen, >1 = overshoot |
@@ -351,19 +354,17 @@ must match `--dit`. The merge is written into a cached copy of the DiT under
351
  on repeat runs, so the ~5–15 s patch cost is paid once. A medium fp32 cache entry
352
  is ~5.4 GB β€” delete `lora_cache/` to reclaim.
353
 
354
- Requires an un-quantized DiT: `fp32` (any family) or `w16a32` (small DiTs only).
355
- Medium now defaults to `w8a8`, but when you pass `--lora` without pinning a
356
- precision the CLI **auto-selects `fp32`** for you (it prints a one-line note); you
357
- only hit the error below if you *explicitly* pin an int8 precision alongside
358
- `--lora`. **LoRA on the quantized
359
- DiTs (`w8a32` / `w8a8-dyn` / `w4a32`) isn't figured out yet** β€” those store
360
- weights as GPTQ-calibrated int8, so merging would mean dequantize β†’ add the LoRA
361
- delta β†’ re-quantize, and a naive re-quant throws away the GPTQ error-feedback
362
- grid (you'd get a model that's both LoRA-adapted *and* degraded). Doing it well
363
- needs a GPTQ pass over the merged weights; until then `--lora` refuses quantized
364
- precisions (the CLI errors, the web UI disables the LoRA panel). Use fp32/w16a32
365
- for LoRA β€” on CPU fp32 is the fast-and-accurate choice anyway, so this is rarely
366
- a real constraint. **Per-step gating (`steps=`) is also MLX-only** β€” a frozen
367
  TFLite graph merges weights once at load; use `optimized/mlx` for step-gated LoRA.
368
 
369
  ## Files
@@ -386,10 +387,12 @@ sa3_tflite/
386
  β”œβ”€β”€ tokenizer.model ← SentencePiece model, BUNDLED (~4 MB; T5Gemma tflite is encoder-only)
387
  β”œβ”€β”€ defs/
388
  β”‚ └── tflite_pipeline.py ← Tokenizer + T5Gemma front-end + pingpong schedule + sampler + WAV
389
- └── tflite/ ← .tflite models (lazy-downloaded; default bundle ~2.4 GB per DiT; ~9.5 GB if you fetch every tier)
390
  β”œβ”€β”€ t5gemma/encoder_fp16.tflite 564 MB text encoder (fp16)
391
- β”œβ”€β”€ sa3-sm-music/dit_fp32.tflite 1.8 GB small music DiT (conditioner baked in)
392
- β”œβ”€β”€ sa3-sm-sfx/dit_fp32.tflite 1.8 GB small sfx DiT (conditioner baked in)
 
 
393
  β”œβ”€β”€ sa3-m/dit_w8a8.tflite 1.6 GB medium DiT β€” DEFAULT (cache-safe rung int8)
394
  β”œβ”€β”€ sa3-m/dit_fp32.tflite 5.8 GB medium DiT (fp32 reference; LoRA / A-B)
395
  β”œβ”€β”€ same-s/{enc,dec}_{fp32,w8a8}.tflite rung codec (fp32 ~0.22 GB Β· w8a8 ~0.07 GB); legacy/ = pre-rung
@@ -413,11 +416,12 @@ is the one weight that IS committed, since the `.tflite` T5Gemma is encoder-only
413
 
414
  ## Notes on the design
415
 
416
- - **Baked-I/O varlen graphs.** Each `.tflite` is a single self-contained graph
417
- with the conditioner and patch/unpatch in-graph, accepting a variable sequence
418
- length β€” so one file serves any `--seconds`. The DiT is a 6-input graph
419
- (`x, t, t5_hidden, t5_mask, seconds, local_add_cond`); feed raw T5 outputs and
420
- the in-graph conditioner handles prompt-padding + seconds-embedding.
 
421
  - **Per-family precision (no longer fp32-only).** The medium DiT now ships the
422
  cache-safe **rung `w8a8`** as its default: int8 weights with per-token dynamic
423
  activation quant, so on an int8-capable CPU (VNNI/AMX) the int8Γ—int8 matmuls run
@@ -425,10 +429,9 @@ is the one weight that IS committed, since the `.tflite` T5Gemma is encoder-only
425
  int8 on the DiT is a *different* (not worse) sample β€” per-step error compounds over
426
  the 8 chaotic sampling steps into another plausible draw, not a degraded one β€” so the
427
  bit-exact `fp32` tier ships alongside it (for CPUs without int8 acceleration, LoRA
428
- merges, and A/B reference). fp16 (`w16a32`) is *slower* on CPU β€” XNNPACK dequantizes it
429
- to fp32 to matmul β€” so it's offered only as a half-size near-lossless tier, not a speed
430
- one. The small DiTs stay fp32 until they're rung-built. T5Gemma is fp16 (numerically
431
- lossless there, and it halves that file).
432
  - **Monotonic audio-to-audio schedule.** The pingpong schedule applies the LogSNR
433
  shift to the normalized `[1β†’0]` grid, then scales by Οƒmax, so audio-to-audio
434
  (Οƒmax < 1) stays monotonically decreasing while keeping all N distilled steps.
@@ -440,13 +443,10 @@ is the one weight that IS committed, since the `.tflite` T5Gemma is encoder-only
440
  RAM flat. This replaced the dynamic-varlen graphs whose RAM blew up (SAME-L O(SΒ²) β†’
441
  16 GB@256; SAME-S β†’ ~50 GB@8192). See [docs/RUNGS.md](docs/RUNGS.md).
442
  - **CFG (`--cfg β‰  1`)** combines a cond and an uncond velocity in denoised space
443
- (optional APG). The canonical DiT is **variable-batch**, so by default cond+uncond
444
- run as **one batch=2 invoke per step** (`--cfg-batched`) β€” ~7–29% faster on
445
- Apple-Silicon (the AMX matrix unit amortizes the weight loads across both rows;
446
- measured on an M4 Pro MacBook Pro).
447
- `--no-cfg-batched` falls back to a sequential batch=1 dual-pass (like the TensorRT
448
- release, whose engine is static-batch=1); the two are bit-identical at `fp32`/`w8a32`
449
- (~80 dB at `w16a32`; `w8a8-dyn` diverges by design β€” batch-shared activation scales).
450
 
451
  ## License & attribution
452
 
 
63
 
64
  Precision is split between the **DiT** and the **SAME codec** β€” they react to quantization oppositely:
65
 
66
+ - **DiT** (`--dit-precision`, or the global `--precision`) β€” **all three DiTs** (medium, sm-music, sm-sfx) are
67
+ from-repo **rung ladders** (see *The rung-ladder DiTs* below), each with the same two tiers:
68
+ - **`w8a8` (default)** β€” **cache-safe int8**: ~equal quality to the old `w8a8-dyn` (cos β‰₯ 0.996 per forward)
69
+ but with **flat RAM**, faster at the rung sizes, and a single self-contained file.
70
+ - **`fp32`** β€” the bit-exact reference (and the pick for LoRA merges or an A/B).
71
+
72
+ int8 on the DiT is a *different* (not worse) sample β€” the 8-step sampler is chaotically sensitive β€” so `w8a8`
73
+ is the shipped default and `fp32` is the reference. medium's ladder tops at 4096 latents (380 s); the small DiTs'
74
+ ladder tops at 1292 (their 2-min max). The old varlen tiers (`w16a32` / `w8a32` / `w8a8-dyn`) are retired to
75
+ `tflite/sa3-*/legacy/`. (wXaY = weight/activation bits; "16" = fp16.)
76
  - **Codec** (`--decoder-precision` / `--encoder-precision`): `fp32` Β· **`w8a8` (default)**. The SAME enc/dec ship as
77
  static **rung** models β€” see [docs/RUNGS.md](docs/RUNGS.md). They run once on a fixed latent, so int8 is
78
  quality-free: a real-music round-trip is identical to fp32. `w8a8` is ~3Γ— faster and ~half the RAM on an
 
80
  codec variants added no audible quality and are retired to `tflite/same-*/legacy/`.)
81
 
82
  `--precision` is a global default that sets the DiT directly and maps to the codec (`fp32`β†’`fp32`, any int8β†’`w8a8`).
83
+ When omitted, the DiT is **`w8a8` for every family** and the **codec is `w8a8`**. All variants keep the full feature
84
+ set (any `--seconds`, odd/even L) and lazy-download from HuggingFace on first use.
 
85
 
86
  | component | precisions | size | notes |
87
  |---|---|---|---|
88
+ | **DiT β€” medium** (rung) | **`w8a8`** (default) / `fp32` | 1.6 / 5.8 GB | 12-rung ladder (β†’ 4096 / 380 s); `w8a8` = cache-safe int8, **flat ~2–4 GB RAM**, β‰ˆ old `w8a8-dyn`; old `w16a32/w8a32/w8a8-dyn` β†’ `legacy/`; litert β‰₯ 2.2.0 |
89
+ | **DiT β€” small** (rung) | **`w8a8`** (default) / `fp32` | 0.5 / 1.8 GB | 8-rung ladder (β†’ 1292 / 2 min); same cache-safe int8, **flat ~0.8 GB RAM** (4.5Γ— less than the old varlen at 2 min); old tiers β†’ `legacy/` |
90
  | **codec** (rung) | `fp32` / `w8a8` | SAME-S enc/dec 0.22 / 0.25 β†’ w8a8 0.06 / 0.09 GB Β· SAME-L 1.7 / 1.8 β†’ w8a8 0.47 / 0.52 GB | w8a8 quality-free, ~3Γ— faster + half RAM; needs litert β‰₯ 2.2.0 |
91
  | T5Gemma | fp16 | 0.6 GB | single-precision |
92
 
93
+ *(The rung `w8a8` uses per-channel int8 weights + per-token dynamic activation quant, so its int8Γ—int8 matmuls
94
+ run faster than fp32 on an int8-capable CPU. The codec rungs use dynamic int8 too, where the extra activation-quant
95
+ error is far below the AE's own loss.)*
96
 
97
+ All three rung DiTs run **batch=1 / sequential CFG** (one forward at `--cfg 1`, a cond+uncond dual-pass otherwise),
98
+ so `--cfg-batched` is a no-op β€” the old varlen "batched CFG" behavior retired with the varlen graphs.
 
 
99
 
100
+ ### The rung-ladder DiTs (replace the old varlen tiers)
101
 
102
+ Every DiT β€” medium **and** the two small models β€” now ships as a **static rung ladder**, built from the repo's own
103
+ PyTorch model (see [build/](build/README.md)), not a dynamic-shape varlen graph. A render of length *L* runs on the
104
+ **smallest fixed rung β‰₯ L** in one forward, with the extra positions masked out of self-attention β€” so the result is
105
+ **exact, not tiled**. Both `fp32` and the cache-safe int8 `w8a8` are single self-contained files (`dit_{fp32,w8a8}.tflite`;
106
+ the per-step global-cond preamble rides inside as a `gcond` signature, so there is no extra file to ship).
107
 
108
  **Why it matters β€” RAM.** The old varlen graph materialized the full `[heads, S, S]` attention, so peak RAM grew with
109
  length and hit **~40 GB at the 380 s maximum** β€” unusable on a normal machine. The rung ladder's static shapes let
 
113
 
114
  ![Medium DiT peak RAM β€” rung ladder vs varlen](docs/img/dit-ram-rung-vs-varlen.png)
115
 
116
+ **Medium β€” 12 rungs** (in latents; audio β‰ˆ L / 10.8 s, 4096 = 380 s max):
117
+ `8, 48, 96, 192, 416, 704, 1056, 1496, 2040, 2824, 3536, 4096` (via `DIT_LADDER` in [`build/build_dit.sh`](build/build_dit.sh)).
 
 
 
118
 
119
+ **Small (sm-music / sm-sfx) β€” 8 rungs**, capped at the 2-minute maximum (`valid_T_lat(120) = 1292` latents):
120
+ `8, 48, 96, 192, 416, 704, 1056, 1292` (built by [`build/build_dit_small.sh`](build/build_dit_small.sh)). Same win β€”
121
+ peak RAM stays **flat at ~0.8 GB** across the whole range, vs the old varlen graph's quadratic climb to ~3.8 GB at
122
+ 2 min (**4.5Γ— less**). sm-music and sm-sfx are architecturally identical, so the curve is the same for both:
123
+
124
+ ![Small DiT peak RAM β€” rung ladder vs varlen](docs/img/dit-ram-rung-vs-varlen-small.png)
125
+
126
+ The short rungs keep tiny clips fast (a 1 s clip no longer pays for a long render); a length *between* rungs pads up to
127
+ the next one. The old varlen tiers (`fp32`, `w16a32`, `w8a32`, `w8a8-dyn`) for all three families are retired to
128
+ **`tflite/sa3-{m,sm-music,sm-sfx}/legacy/`** on HuggingFace.
129
 
130
  ```bash
131
  # medium DiT β€” defaults: w8a8 rung DiT + w8a8 SAME-S codec (the `--decoder` now defaults to same-s)
 
150
  you get; the encoder's affects a2a/inpaint latents):
151
 
152
  ```bash
153
+ # fp32 DiT (bit-exact reference), reference-quality fp32 codec
154
  ./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s \
155
+ --dit-precision fp32 --decoder-precision fp32
156
  ```
157
 
158
  ## Install
 
268
  first run it offers to install the UI-only extras (`gradio`, `pillow`,
269
  `soundfile`) into `.venv`. In the browser:
270
 
271
+ - **Precision** dropdown next to the model picker β€” `w8a8` (default) / `fp32`,
272
+ applied to the DiT; the codec maps it to its rung tier (`fp32` β†’ fp32, int8 β†’ w8a8).
 
273
  - **LoRA** panel β€” upload a `.safetensors` adapter or pick one from
274
  `loras/<model>/` (e.g. `loras/sa3-medium/`), set its strength, stack several;
275
  settings are remembered per model. The panel is disabled under a quantized
 
313
  | `--seed` | random | Set for reproducibility; the chosen seed is printed at the end |
314
  | `--cfg` | 1.0 | Guidance scale; 1.0 = off, >1 toward prompt, <1 toward uncond. β‰ 1 runs cond+uncond each step |
315
  | `--apg` | 1.0 | Adaptive Projected Guidance; only matters when `--cfg β‰  1` |
316
+ | `--cfg-batched` | on | No-op on the rung DiTs (all batch=1): when `--cfg β‰  1` they always run a sequential cond+uncond dual-pass. Kept for interface parity with the MLX/TRT backends |
317
+ | `--lora` | β€” | A `.safetensors` LoRA adapter (SA3-native/underfit or PEFT) merged into the DiT, optional `strength=S`; repeat to stack. Requires `--dit-precision fp32` (the int8 w8a8 rung can't be merged). See [LoRA](#lora) |
318
  | `--lora-strength` | 1.0 | Default strength for `--lora` adapters without their own `strength=` |
319
  | `--init-audio` | β€” | WAV (any format via ffmpeg) input for audio-to-audio / inpaint |
320
  | `--init-noise-level` | 1.0 | Οƒmax; 0.4–0.8 typical for variation, 1.0 = full regen, >1 = overshoot |
 
354
  on repeat runs, so the ~5–15 s patch cost is paid once. A medium fp32 cache entry
355
  is ~5.4 GB β€” delete `lora_cache/` to reclaim.
356
 
357
+ Requires the `fp32` DiT. Every family now defaults to `w8a8`, but when you pass
358
+ `--lora` without pinning a precision the CLI **auto-selects `fp32`** for you (it
359
+ prints a one-line note); you only hit the error below if you *explicitly* pin
360
+ `w8a8` alongside `--lora`. **LoRA on the int8 `w8a8` rung isn't figured out yet** β€”
361
+ it stores weights as calibrated int8, so merging would mean dequantize β†’ add the
362
+ LoRA delta β†’ re-quantize, and a naive re-quant throws away the quantization grid
363
+ (you'd get a model that's both LoRA-adapted *and* degraded). Doing it well needs a
364
+ quant pass over the merged weights; until then `--lora` refuses `w8a8` (the CLI
365
+ errors, the web UI disables the LoRA panel). Use `fp32` for LoRA β€” on CPU it's the
366
+ fast-and-accurate choice anyway, so this is rarely a real constraint. **Per-step
367
+ gating (`steps=`) is also MLX-only** β€” a frozen
 
 
368
  TFLite graph merges weights once at load; use `optimized/mlx` for step-gated LoRA.
369
 
370
  ## Files
 
387
  β”œβ”€β”€ tokenizer.model ← SentencePiece model, BUNDLED (~4 MB; T5Gemma tflite is encoder-only)
388
  β”œβ”€β”€ defs/
389
  β”‚ └── tflite_pipeline.py ← Tokenizer + T5Gemma front-end + pingpong schedule + sampler + WAV
390
+ └── tflite/ ← .tflite models (lazy-downloaded; default bundle ~0.7 GB small / ~2.4 GB medium)
391
  β”œβ”€β”€ t5gemma/encoder_fp16.tflite 564 MB text encoder (fp16)
392
+ β”œβ”€β”€ sa3-sm-music/dit_w8a8.tflite 0.5 GB small music DiT β€” DEFAULT (cache-safe rung int8)
393
+ β”œβ”€β”€ sa3-sm-music/dit_fp32.tflite 1.8 GB small music DiT (fp32 reference; LoRA / A-B)
394
+ β”œβ”€β”€ sa3-sm-sfx/dit_w8a8.tflite 0.5 GB small sfx DiT β€” DEFAULT (cache-safe rung int8)
395
+ β”œβ”€β”€ sa3-sm-sfx/dit_fp32.tflite 1.8 GB small sfx DiT (fp32 reference; LoRA / A-B)
396
  β”œβ”€β”€ sa3-m/dit_w8a8.tflite 1.6 GB medium DiT β€” DEFAULT (cache-safe rung int8)
397
  β”œβ”€β”€ sa3-m/dit_fp32.tflite 5.8 GB medium DiT (fp32 reference; LoRA / A-B)
398
  β”œβ”€β”€ same-s/{enc,dec}_{fp32,w8a8}.tflite rung codec (fp32 ~0.22 GB Β· w8a8 ~0.07 GB); legacy/ = pre-rung
 
416
 
417
  ## Notes on the design
418
 
419
+ - **Baked-I/O rung graphs.** Each `.tflite` is a single self-contained graph with the
420
+ conditioner and patch/unpatch in-graph, so you feed raw T5 outputs and the graph handles
421
+ prompt-padding + seconds-embedding. One file serves any `--seconds` by dispatching to the
422
+ **smallest fixed rung β‰₯ L** (extra positions masked out of attention β€” exact, not tiled).
423
+ The DiT rung takes `x, gc, t5_hidden, t5_mask, seconds, local_add_cond, attn_mask`; the
424
+ per-step global-cond `gc` is produced by the bundled `gcond` signature in the same file.
425
  - **Per-family precision (no longer fp32-only).** The medium DiT now ships the
426
  cache-safe **rung `w8a8`** as its default: int8 weights with per-token dynamic
427
  activation quant, so on an int8-capable CPU (VNNI/AMX) the int8Γ—int8 matmuls run
 
429
  int8 on the DiT is a *different* (not worse) sample β€” per-step error compounds over
430
  the 8 chaotic sampling steps into another plausible draw, not a degraded one β€” so the
431
  bit-exact `fp32` tier ships alongside it (for CPUs without int8 acceleration, LoRA
432
+ merges, and A/B reference). **All three families** (medium + both small) now ship this
433
+ rung `fp32`/`w8a8` pair; the old fp16/int8 varlen tiers (`w16a32`/`w8a32`/`w8a8-dyn`) are
434
+ retired to `legacy/`. T5Gemma is fp16 (numerically lossless there, and it halves that file).
 
435
  - **Monotonic audio-to-audio schedule.** The pingpong schedule applies the LogSNR
436
  shift to the normalized `[1β†’0]` grid, then scales by Οƒmax, so audio-to-audio
437
  (Οƒmax < 1) stays monotonically decreasing while keeping all N distilled steps.
 
443
  RAM flat. This replaced the dynamic-varlen graphs whose RAM blew up (SAME-L O(SΒ²) β†’
444
  16 GB@256; SAME-S β†’ ~50 GB@8192). See [docs/RUNGS.md](docs/RUNGS.md).
445
  - **CFG (`--cfg β‰  1`)** combines a cond and an uncond velocity in denoised space
446
+ (optional APG). The rung DiTs are **static batch=1**, so cond+uncond always run as a
447
+ **sequential dual-pass** (like the TensorRT release, whose engine is static-batch=1);
448
+ `--cfg-batched` is a no-op, kept only for interface parity with the MLX/TRT backends.
449
+ (The old variable-batch "batched CFG" fast-path retired with the varlen graphs.)
 
 
 
450
 
451
  ## License & attribution
452
 
tflite/docs/img/dit-ram-rung-vs-varlen-small.png ADDED
tflite/sa3-sm-music/dit_fp32.tflite CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:d388700a2ca439c11e9a53506e964e93231386a2beb8173c6eec6d95f676ce09
3
- size 1838758544
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3fd214377c48f31521aeee08fd1202ab0a8ea4135d2358f6cec9a443f7bf35f6
3
+ size 1845785264
tflite/sa3-sm-music/dit_w8a8.tflite ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1d8ef99bf10f51cc52b68f7b527d3773828e700758b253edd8a8e088a906e241
3
+ size 539350360
tflite/sa3-sm-music/legacy/dit_fp32.tflite ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d388700a2ca439c11e9a53506e964e93231386a2beb8173c6eec6d95f676ce09
3
+ size 1838758544
tflite/sa3-sm-music/{dit_w16a32.tflite β†’ legacy/dit_w16a32.tflite} RENAMED
File without changes
tflite/sa3-sm-music/{dit_w8a32.tflite β†’ legacy/dit_w8a32.tflite} RENAMED
File without changes
tflite/sa3-sm-music/{dit_w8a8-dyn.tflite β†’ legacy/dit_w8a8-dyn.tflite} RENAMED
File without changes
tflite/sa3-sm-sfx/dit_fp32.tflite CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:6060ecfeca34c4ab35bc1912a37e680e8cd7aab6c4bd9de1bc2655414891b8d8
3
- size 1838758544
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:06ea22d57bd94318ea4e11d5d4c104f2bb17eca0540623db048bdef33aba5753
3
+ size 1845785264
tflite/sa3-sm-sfx/dit_w8a8.tflite ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d3f99409f446eff9a625768ba2ee09a2347ea90c4897c1662fbcd8837fae3045
3
+ size 539350360
tflite/sa3-sm-sfx/legacy/dit_fp32.tflite ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6060ecfeca34c4ab35bc1912a37e680e8cd7aab6c4bd9de1bc2655414891b8d8
3
+ size 1838758544
tflite/sa3-sm-sfx/{dit_w16a32.tflite β†’ legacy/dit_w16a32.tflite} RENAMED
File without changes
tflite/sa3-sm-sfx/{dit_w8a32.tflite β†’ legacy/dit_w8a32.tflite} RENAMED
File without changes
tflite/sa3-sm-sfx/{dit_w8a8-dyn.tflite β†’ legacy/dit_w8a8-dyn.tflite} RENAMED
File without changes