GraftMuLa

HeartCodec RVQ token-level style transfer

Installation

!pip install -U graftmula

Introduction and usage

HeartCodec RVQ token surgery: content/texture level swapping between two tracks.

Tier-C style transfer on the HeartCodec token space. Provides BOTH a CLI and a programmatic API (TokenSurgeryConfig / TokenSurgery) with identical options.

Idea

HeartCodec encodes audio into 8 residual-vector-quantization (RVQ) token streams at 12.5 Hz. The low indices carry the coarse, semantic/musical content (level 0 is deliberately distilled from semantic features) while the high indices carry fine acoustic detail (texture, transients, noise, production sheen). This script tokenizes two tracks:

• source the CONTENT donor (defines output length, melody, structure)
• style the TEXTURE donor (defines the sonic surface)

then builds a hybrid and decodes it with the flow-matching decoder:

• levels 0 .. C-1 come from the source (content_levels)
• levels C .. 7 come from the style track (optionally offset/tiled)

Surgery modes

mode="swap" (default) Exact token-ID transplant. Levels are copied verbatim and decoded through the stock HeartCodec.detokenize. Fully reversible, no quantizer internals touched.

mode="blend" Embedding-space mix. Both tracks are dequantized to per-level codebook embeddings (project_out of each level's code); the hybrid conditioning is the per-level weighted sum (source weight per level, style weight = 1-w), applied with a SINGLE project_out so the bias term matches the training path exactly. v2 decodes this embedding DIRECTLY -- no re-quantization.

v1 re-quantized the mix into token IDs before decoding. That nearest-
codebook snap re-derives ALL levels from scratch and, for fractional
weights, lands on temporally incoherent token tuples -> garbled audio
(amplified by guidance > 1.25). Direct embedding decode eliminates the
snap: the conditioning is exactly the interpolated point. Set
``requantize=True`` only if you explicitly want a best-effort token-tuple
representation (e.g. to dump tokens for a future style prior).

Decoder internals used by blend mode (verified against heartlib commit 5858ca8 / vector-quantize-pytorch==1.27.15): HeartCodec.tokenize(waveform, sr) -> int64 [8, T] on CPU, 12.5 Hz HeartCodec.detokenize(codes, ...) -> stereo [2, ~T/12.5*48k] HeartCodecDecoder.detokenize -> windowed flow decode, window duration 29.76 s, token hop 320 frames, 104-latent-frame context stitching, overlap-add crossfade (transcribed verbatim in detokenize_embeddings with the code lookup replaced by an injected embedding sequence) ResidualVQ.get_codes_from_indices([B,T,Q]) -> [Q, B, T, codebook_dim] ResidualVQ.project_out -> nn.Linear(codebook_dim, 512) ResidualVQ.forward([B,T,512]) -> (quant, indices [B,T,Q], loss)

Because the embedding path is a transcription, verify_embed_decode=True (default in blend mode) cross-checks it against the stock decoder on a single window with identical RNG and prints TWO numbers: a conditioning-space diff (isolates quantizer-path drift) and a waveform diff. On a matched environment (heartlib commit 5858ca8, vq-pytorch==1.27.15) both are 0.0. A small nonzero waveform diff signals heartlib/vq-pytorch version or CUDA-numerics drift and only WARNS; the run aborts on non-finite or > 0.5 diffs, or at the 5e-2 threshold when strict_parity=True.

Performance options

torch_compile + compile_mode Wraps flow_matching.estimator (the DiT, called 2 x num_steps x windows times) and scalar_model.decode in torch.compile. First decode pays an inductor warm-up. "reduce-overhead" enables CUDA graphs: heartlib's euler loop keeps tensors alive across invocations, which can trip "accessing tensor output of CUDAGraphs that has been overwritten" errors. The module mitigates this automatically (per-invocation step marks plus inductor generation cloning) and, if a CUDA-graph error still fires, falls back to eager mid-run and stays eager. "default" and "max-autotune-no-cudagraphs" never use CUDA graphs.

amp_autocast + amp_dtype Wraps decoding in torch.autocast. The checkpoint stays FP32; autocast only casts matmul/conv inputs. bfloat16 is the safe default; float16 can overflow in attention. Tokenization always stays FP32.

Programmatic usage

from graftmula import TokenSurgery, TokenSurgeryConfig

ts = TokenSurgery(TokenSurgeryConfig(mode="blend", content_levels=3,
                                     ramp=2, guidance_scale=3.0))
src = ts.tokenize_file("content.wav")
sty = ts.tokenize_file("texture.wav")
hybrid = ts.build_hybrid(src, sty)            # HybridTokens
wav = ts.decode(hybrid, seed=42)              # [2, S] cpu float32
# or one shot (writes files, returns SurgeryResult):
result = ts.run("content.wav", "texture.wav", "hybrid.wav")
# or sweep cut points: every (content_levels x ramp) variant from ONE
# tokenization pass (files <stem>.cl{C}[_r{R}].<ext>, returns GridResult):
grid = ts.compare_grid("content.wav", "texture.wav", "grid.wav",
                       levels=(1, 2, 3, 4, 5), ramps=(0, 2))

CLI usage

python graftmula.py \
    --decoder_path ./ckpt/HeartCodec-oss-20260123 \
    --encoder_path ./ckpt/HeartCodec-oss-encoder \
    --source_path ./assets/source_content.mp3 \
    --style_path  ./assets/style_texture.mp3 \
    --mode blend --content_levels 3 --ramp 2 \
    --output_path ./assets/hybrid.wav

# A/B grid: decode every cut point in one pass (tokens tokenized once):
python graftmula.py \
    --decoder_path ./ckpt/HeartCodec-oss-20260123 \
    --encoder_path ./ckpt/HeartCodec-oss-encoder \
    --source_path ./assets/source_content.mp3 \
    --style_path  ./assets/style_texture.mp3 \
    --mode blend --compare_grid --grid_levels 1,2,3,4,5 --grid_ramps 0,2 \
    --output_path ./assets/grid.wav

Practical tips

• Keep seed fixed while sweeping content_levels; this module reseeds before every decode, so differences come from the surgery alone.
• Content levels 1-4 are the interesting sweep range. 0 = fully restyled, 8 = no-op.
• Style texture tiles if the style segment is shorter than the source; pick a dense passage with style_offset.
• ramp weights descend in uniform steps from the cut and may truncate below full-style at level 7 (only 8 levels exist). Inspect the printed weights or set level_weights explicitly.
• The inner solver progress bar (inside solve_euler) is always shown by heartlib; quiet only suppresses our own logging (kept for parity).
• --compare_grid renders every (content_levels x ramp) variant from a single tokenization pass -- the cheap way to find the right cut point. The parity check and the source/style baselines run once, not per variant, so each extra grid cell costs exactly one decode.

Attribution

GraftMuLa was vibe-coded by Z AI GLM 5.3 Flash AI Model exlusively for HeartMuLa project:

https://chat.z.ai/ https://github.com/HeartMuLa https://github.com/HeartMuLa/heartlib

License

Apache 2.0

Project Los Angeles

Tegridy Code 2026

Project Los Angeles

Tegridy Code 2026