"""Complete Audio8 encoder/projector/conditioned Qwen2 decoder and tied head.""" from __future__ import annotations import json from pathlib import Path import mlx.core as mx from .weights import Weights, require, sha256 from .math import KVCache, attention, gelu, rms_norm, fast_rms_norm, rope, rotary_factors, apply_rotary DTYPES = {'float32': mx.float32, 'float16': mx.float16, 'bfloat16': mx.bfloat16} class Audio8Model: def __init__(self, weights, config, frontend, *, dtype='float32', cache_dtype='bfloat16', fuse_projections=False, math_mode='reference'): require(config['text_config']['model_type'] == 'qwen2', 'only pinned Qwen2 text architecture') require(config['max_frame_len'] == 8 and config['frame_lens'] == [4, 6, 8], 'projector/frame configuration') self.weights, self.config, self.frontend = weights, config, frontend self.dtype, self.cache_dtype = DTYPES[dtype], DTYPES[cache_dtype] self.dtype_name, self.cache_dtype_name = dtype, cache_dtype require(math_mode in ('reference', 'shared-rope', 'compiled'), 'math mode') self.math_mode = math_mode self.compiled_blocks = {} self.audio = config['audio_config']; self.text = config['text_config'] self.audio_theta = self.audio['rope_parameters']['rope_theta'] self.text_theta = self.text['rope_parameters']['rope_theta'] for i in range(self.audio['num_hidden_layers']): self._validate_layer(f'audio_tower.layers.{i}', True) for i in range(self.text['num_hidden_layers']): self._validate_layer(f'language_model.model.layers.{i}', False) self.fused_projection_groups = 0 if fuse_projections: for stem, layers in [('audio_tower.layers', self.audio['num_hidden_layers']), ('language_model.model.layers', self.text['num_hidden_layers'])]: for index in range(layers): p = f'{stem}.{index}' for suffixes in [('self_attn.q_proj', 'self_attn.k_proj', 'self_attn.v_proj'), ('mlp.gate_proj', 'mlp.up_proj')]: self.fused_projection_groups += weights.fuse([p + '.' + suffix for suffix in suffixes]) head = weights.tensors['language_model.model.embed_tokens.weight'] require((head.rows, head.cols) == (self.text['vocab_size'], self.text['hidden_size']), 'embedding dimensions') def _validate_layer(self, prefix, encoder): config = self.audio if encoder else self.text h, n, intermediate = config['hidden_size'], config['num_attention_heads'], config['intermediate_size'] dim = config.get('head_dim', h // n) kv = config.get('num_key_value_heads', n) for sub, shape in {'self_attn.q_proj': (n * dim, h), 'self_attn.k_proj': (kv * dim, h), 'self_attn.v_proj': (kv * dim, h), 'self_attn.o_proj': (h, n * dim), 'mlp.gate_proj': (intermediate, h), 'mlp.up_proj': (intermediate, h), 'mlp.down_proj': (h, intermediate)}.items(): weight = self.weights.tensors[prefix + '.' + sub + '.weight'] require((weight.rows, weight.cols) == shape, f'layer weight shape: {prefix}.{sub}') @classmethod def load(cls, bundle, config_path, frontend, **kwargs): config = json.loads(Path(config_path).read_text()) weights = Weights.load(bundle) weights.provenance['config_sha256'] = sha256(config_path) require(weights.provenance['expected_config_sha256'] == weights.provenance['config_sha256'], 'config differs from bundle source asset') return cls(weights, config, frontend, **kwargs) def session(self, *, gear=4, delay_tokens=3, context=375, trim=38, stable=16): return Session(self, gear, delay_tokens, context, trim, stable) def warm_compiled_projections(self, prefill, batch_windows): """Compile bounded projection shapes before audio readiness. Only pure projection functions receive synthetic inputs. No session, tokenizer, audio reader, position, or KV cache is created or advanced. Explicit weights are shared with subsequent inference unchanged. """ require(self.math_mode == 'compiled', 'projection warmup requires compiled math') require(type(prefill) is int and 1 <= prefill <= 64, 'warmup prefill') require(type(batch_windows) is int and batch_windows in (1, 2, 4), 'warmup batch') calls = 0 for encoder, config, stem, lengths in ( (True, self.audio, 'audio_tower.layers', sorted({4 * prefill, *range(4, 4 * batch_windows + 1, 4)})), (False, self.text, 'language_model.model.layers', sorted({1, prefill})), ): modulation = (mx.array(1, self.dtype) if encoder else mx.ones((1, config['hidden_size']), self.dtype)) for index in range(config['num_hidden_layers']): (qkv, qp), (ffn, fp) = self._projection_blocks(f'{stem}.{index}', encoder) for length in lengths: x = mx.zeros((1, length, config['hidden_size']), self.dtype) mx.eval(*qkv(x, qp), ffn(x, modulation, fp)) calls += 2 mx.synchronize() return calls def _projection_blocks(self, prefix, encoder): """Pure stateless functions: no cache/position mutation or array copies. Hidden dimensions, weight identity, epsilon and norm dtype are fixed. Two function objects per layer cache only bounded token-length shapes: at most the initial encoder length plus4/8/12/16, and initial/one-token decoder lengths. No position/history length enters these functions. """ if prefix not in self.compiled_blocks: cfg = self.audio if encoder else self.text norm1 = 'self_attn_layer_norm' if encoder else 'input_layernorm' norm2 = 'final_layer_norm' if encoder else 'post_attention_layernorm' weights, eps = self.weights, cfg['rms_norm_eps'] norm1_weight = weights.tensor(prefix + '.' + norm1 + '.weight') norm2_weight = weights.tensor(prefix + '.' + norm2 + '.weight') qkv_names = [prefix + '.self_attn.' + name for name in ('q_proj', 'k_proj', 'v_proj')] gate_up_names = [prefix + '.mlp.gate_proj', prefix + '.mlp.up_proj'] qkv_op, qkv_arrays = weights.linear_plan(qkv_names) gate_up_op, gate_up_arrays = weights.linear_plan(gate_up_names) down_op, down_arrays = weights.linear_plan([prefix + '.mlp.down_proj']) def qkv(x, parameters): z = fast_rms_norm(x, parameters['norm'], eps) return qkv_op(z, parameters['projection']) def ffn(x, modulation, parameters): z = fast_rms_norm(x, parameters['norm'], eps) * modulation gate, up = gate_up_op(z, parameters['gate_up']) return x + down_op((gate * mx.sigmoid(gate)) * up, parameters['down'])[0] # Explicit function arguments, rather than closure constants or # a mutable captured list, preserve exact shared weight ownership. qkv_inputs = {'norm': norm1_weight, 'projection': qkv_arrays} ffn_inputs = {'norm': norm2_weight, 'gate_up': gate_up_arrays, 'down': down_arrays} # MLX0.32 cannot infer fused-output slice shapes in shapeless mode # when packed arrays are explicit arguments. Bounded fixed shapes # are valid here; unlike whole layers, no growing KV shape enters. self.compiled_blocks[prefix] = ((mx.compile(qkv), qkv_inputs), (mx.compile(ffn), ffn_inputs)) return self.compiled_blocks[prefix] def layer(self, x, cache, prefix, position, encoder, modulation=None, attention_partitions=None, rotary=None): cfg = self.audio if encoder else self.text heads = cfg['num_attention_heads']; dim = cfg.get('head_dim', cfg['hidden_size'] // heads) kv_heads = cfg.get('num_key_value_heads', heads) norm1 = 'self_attn_layer_norm' if encoder else 'input_layernorm' norm2 = 'final_layer_norm' if encoder else 'post_attention_layernorm' w = self.weights compiled = getattr(self, 'math_mode', 'reference') == 'compiled' if compiled: block, parameters = self._projection_blocks(prefix, encoder)[0] q, k, v = block(x, parameters) else: z = rms_norm(x, w.tensor(prefix + '.' + norm1 + '.weight'), cfg['rms_norm_eps']) q, k, v = w.linear_many([prefix + '.self_attn.' + name for name in ('q_proj', 'k_proj', 'v_proj')], z) q = q.reshape(1, -1, heads, dim).transpose(0, 2, 1, 3) k = k.reshape(1, -1, kv_heads, dim).transpose(0, 2, 1, 3) v = v.reshape(1, -1, kv_heads, dim).transpose(0, 2, 1, 3) theta = self.audio_theta if encoder else self.text_theta q, k = ((rope(q, position, theta), rope(k, position, theta)) if rotary is None else (apply_rotary(q, rotary), apply_rotary(k, rotary))) if attention_partitions is None: result = attention(q, k, v, cache, window=cfg['sliding_window'] if encoder else None) else: require(encoder and sum(attention_partitions) == q.shape[2] and all(n > 0 for n in attention_partitions), 'attention partitions') contexts, offset = [], 0 for count in attention_partitions: stop = offset + count contexts.append(attention(q[:, :, offset:stop], k[:, :, offset:stop], v[:, :, offset:stop], cache, window=cfg['sliding_window'])) offset = stop result = mx.concatenate(contexts, axis=2) x = x + w.linear(prefix + '.self_attn.o_proj', result.transpose(0, 2, 1, 3).reshape(1, -1, heads * dim)) if compiled: modulation = mx.array(1, x.dtype) if modulation is None else modulation block, parameters = self._projection_blocks(prefix, encoder)[1] return block(x, modulation, parameters) z = rms_norm(x, w.tensor(prefix + '.' + norm2 + '.weight'), cfg['rms_norm_eps']) if modulation is not None: z = z * modulation gate, up = w.linear_many([prefix + '.mlp.gate_proj', prefix + '.mlp.up_proj'], z) return x + w.linear(prefix + '.mlp.down_proj', (gate * mx.sigmoid(gate)) * up) class Session: def __init__(self, model, gear, delay_tokens, context, trim, stable): require(gear in model.config['frame_lens'] and type(delay_tokens) is int and 1 <= delay_tokens <= 30, 'gear/delay') require(context > stable + trim > stable > 1 and trim > 0, 'rolling policy') self.model, self.gear, self.delay_tokens = model, gear, delay_tokens self.context, self.trim, self.stable = context, trim, stable self.position = self.emissions = self.trims = 0 self.encoder = [KVCache(model.audio['sliding_window'] - 1, model.cache_dtype, sliding=True) for _ in range(model.audio['num_hidden_layers'])] self.decoder = [KVCache(context, model.cache_dtype) for _ in range(model.text['num_hidden_layers'])] hidden = model.text['hidden_size'] inv = mx.exp(-mx.log(mx.array(10000., mx.float32)) * mx.arange(hidden // 2, dtype=mx.float32) / (hidden // 2)).astype(model.dtype) phase = mx.array(delay_tokens, model.dtype) * inv condition = mx.concatenate([mx.cos(phase), mx.sin(phase)]) condition = condition + model.weights.tensor('frame_len_embedding.weight')[model.config['frame_lens'].index(gear)].astype(model.dtype) self.modulations = [] for i in range(model.text['num_hidden_layers']): prefix = f'language_model.model.layers.{i}.ada_rms_norm' value = model.weights.linear(prefix + '.linear2', gelu(model.weights.linear(prefix + '.linear1', condition[None]))) self.modulations.append((1 + value).astype(model.dtype)) mx.eval(*self.modulations) def maybe_trim(self, incoming): require(0 < incoming <= self.context - self.stable, 'incoming token count') while self.position + incoming > self.context: require(all(c.length == self.position for c in self.decoder), 'decoder cache clock mismatch') for cache in self.decoder: cache.trim_decoder(self.trim, self.stable, self.model.text_theta) for cache in self.encoder: cache.rebase_encoder(self.trim * self.gear, self.model.audio_theta) # Bound rotation scratch to one layer; preserve the exact operations. mx.eval(*cache.arrays()) self.position -= self.trim; self.trims += 1 mx.eval(*self.cache_arrays()) def cache_arrays(self): return [a for c in self.encoder + self.decoder for a in c.arrays()] @property def cache_bytes(self): return sum(c.nbytes for c in self.encoder + self.decoder) def batch_capacity(self, requested, initial_tokens=1): """Trim before the first query if required, then never cross another trim.""" require(type(requested) is int and 1 <= requested <= 4, 'batch windows must be1..4') require(1 <= initial_tokens <= 64, 'initial token count') position = self.position while position + initial_tokens > self.context: position -= self.trim return min(requested, 1 + self.context - position - initial_tokens) def encode_windows(self, waveforms, counts): m, w = self.model, self.model.weights parts = [m.frontend.convolve(m.frontend.mel(wave), w, count * self.gear, m.dtype) for wave, count in zip(waveforms, counts)] x = mx.concatenate(parts, axis=1) if len(parts) > 1 else parts[0] partitions = [count * self.gear for count in counts] rotary = None if getattr(m, 'math_mode', 'reference') != 'reference': rotary = rotary_factors(self.position * self.gear, x.shape[1], m.audio.get('head_dim', m.audio['hidden_size'] // m.audio['num_attention_heads']), m.dtype, m.audio_theta) for i, cache in enumerate(self.encoder): x = m.layer(x, cache, f'audio_tower.layers.{i}', self.position * self.gear, True, attention_partitions=partitions, rotary=rotary) x = rms_norm(x, w.tensor('audio_tower.norm.weight'), m.audio['rms_norm_eps']) grouped = x.reshape(1, sum(counts), self.gear, m.audio['hidden_size']) grouped = mx.pad(grouped, [(0, 0), (0, 0), (0, 8 - self.gear), (0, 0)]).reshape(1, sum(counts), -1) audio = w.linear('multi_modal_projector.linear_2', gelu(w.linear('multi_modal_projector.linear_1', grouped))) # Finish the encoder batch before autoregressive feedback. Cache history # was rounded at each original window, not only at the batch endpoint. mx.eval(audio, *[a for c in self.encoder for a in c.arrays()]) return audio def decode_audio(self, audio, token_ids, *, return_logits=False): m, w = self.model, self.model.weights x = w.tensors['language_model.model.embed_tokens.weight'].embedding(token_ids, m.dtype)[None] + audio rotary = None if getattr(m, 'math_mode', 'reference') != 'reference': rotary = rotary_factors(self.position, len(token_ids), m.text.get('head_dim', m.text['hidden_size'] // m.text['num_attention_heads']), m.dtype, m.text_theta) for i, cache in enumerate(self.decoder): x = m.layer(x, cache, f'language_model.model.layers.{i}', self.position, False, self.modulations[i], rotary=rotary) x = rms_norm(x[:, -1:], w.tensor('language_model.model.norm.weight'), m.text['rms_norm_eps']) logits = w.tensors['language_model.lm_head.weight'](x)[0, 0].astype(mx.float32) logits = mx.where(mx.arange(logits.size) == m.config['eos_token_id'], -float('inf'), logits) choice = mx.argmax(logits) maximum = mx.max(logits) mx.eval(choice, maximum, *[a for c in self.decoder for a in c.arrays()]) require(bool(mx.isfinite(maximum).item()), 'nonfinite output logits') token = int(choice.item()) self.position += len(token_ids); self.emissions += 1 if return_logits: mx.eval(logits) return (token, logits) if return_logits else token def step_batch(self, waveforms, token_ids, *, return_logits=False, on_token=None): """Batch encoder projections; decode feedback sequentially; no trim inside. At most four source windows. Callback fires after each actual decoder result, so no artificial final-batch timestamp is assigned to early text. """ require(1 <= len(waveforms) <= 4 and 1 <= len(token_ids) <= 64, 'batch shape') self.maybe_trim(len(token_ids)) counts = [len(token_ids)] + [1] * (len(waveforms) - 1) require(self.position + sum(counts) <= self.context, 'batch crosses a rolling boundary') audio = self.encode_windows(waveforms, counts) outputs, offset, inputs = [], 0, token_ids for number, count in enumerate(counts): output = self.decode_audio(audio[:, offset:offset + count], inputs, return_logits=return_logits) outputs.append(output) token = output[0] if isinstance(output, tuple) else output if on_token is not None: on_token(number, output) inputs = [token]; offset += count return outputs def step(self, waveform, token_ids, *, return_logits=False): return self.step_batch([waveform], token_ids, return_logits=return_logits)[0]