""" MLX inference apply_model equivalent. """ from __future__ import annotations import random import typing as tp import warnings import mlx.core as mx from .mlx_utils import center_trim _WEIGHT_CACHE: dict[tuple[int, float, str], mx.array] = {} class TensorChunk: def __init__(self, tensor: mx.array, offset=0, length=None): total_length = tensor.shape[-1] if offset < 0 or offset >= total_length: raise ValueError("Invalid offset.") if length is None: length = total_length - offset else: length = min(total_length - offset, length) if isinstance(tensor, TensorChunk): self.tensor = tensor.tensor self.offset = offset + tensor.offset else: self.tensor = tensor self.offset = offset self.length = length @property def shape(self): shape = list(self.tensor.shape) shape[-1] = self.length return shape def padded(self, target_length): delta = target_length - self.length if delta < 0: raise ValueError("target_length must be >= length.") total_length = self.tensor.shape[-1] start = self.offset - delta // 2 end = start + target_length correct_start = max(0, start) correct_end = min(total_length, end) pad_left = correct_start - start pad_right = end - correct_end out = self.tensor[..., correct_start:correct_end] if pad_left or pad_right: out = mx.pad(out, [(0, 0), (0, 0), (pad_left, pad_right)], mode="constant") return out def tensor_chunk(tensor_or_chunk): if isinstance(tensor_or_chunk, TensorChunk): return tensor_or_chunk if not isinstance(tensor_or_chunk, mx.array): raise TypeError("Expected mx.array.") return TensorChunk(tensor_or_chunk) def apply_model( model, mix: tp.Union[mx.array, "TensorChunk"], shifts: int = 1, split: bool = True, overlap: float = 0.25, transition_power: float = 1.0, progress: bool = False, num_workers: int = 0, segment: tp.Optional[float] = None, batch_size: int = 8, seed: tp.Optional[int] = None, _rng: tp.Optional[random.Random] = None, ): progress_enabled = bool(progress) if num_workers > 0: warnings.warn("num_workers > 0 ignored on MLX.", RuntimeWarning) num_workers = 0 rng: tp.Any if _rng is None: rng = random if seed is None else random.Random(int(seed)) else: rng = _rng # --- BagOfModels Handling --- from .mlx_convert import BagOfModelsMLX if isinstance(model, BagOfModelsMLX): totals = [0.0] * len(model.sources) estimates = None min_length = None for sub_model, model_weights in zip(model.models, model.weights): res = apply_model( sub_model, mix, shifts, split, overlap, transition_power, progress, num_workers, segment, batch_size, seed=seed, _rng=rng ) out = mx.array(res) # Vectorized per-source weighting: (1, S, 1, 1) w = mx.array(model_weights, dtype=out.dtype).reshape(1, -1, 1, 1) out = out * w # Track totals in Python (tiny), keep math vectorized above. for k, inst_weight in enumerate(model_weights): totals[k] += float(inst_weight) if min_length is None: min_length = out.shape[-1] estimates = out else: if out.shape[-1] < min_length: min_length = out.shape[-1] estimates = estimates[..., :min_length] elif out.shape[-1] > min_length: out = out[..., :min_length] estimates = estimates + out # Vectorized normalization by totals. denom = mx.array(totals, dtype=estimates.dtype).reshape(1, -1, 1, 1) estimates = estimates / denom mx.eval(estimates) # Single eval for entire BagOfModels path return estimates # --- Standard Inference --- mix_chunk = tensor_chunk(mix) batch, channels, length = mix_chunk.shape mix_dtype = mix_chunk.tensor.dtype if shifts: max_shift = int(0.5 * model.samplerate) padded_mix = mix_chunk.padded(length + 2 * max_shift) padded_chunk = TensorChunk(padded_mix) out = 0.0 for _ in range(shifts): offset = rng.randint(0, max_shift) shifted = TensorChunk(padded_chunk, offset, length + max_shift - offset) shifted_out = apply_model( model, shifted, 0, split, overlap, transition_power, False, num_workers, segment, batch_size, seed=seed, _rng=rng ) out = out + shifted_out[..., max_shift - offset:] out = out / shifts mx.eval(out) # Single eval after all shifts return out if split: out = mx.zeros((batch, len(model.sources), channels, length), dtype=mix_dtype) sum_weight = mx.zeros((length,), dtype=mix_dtype) if segment is None: segment = model.segment segment_length = int(model.samplerate * segment) stride = int((1 - overlap) * segment_length) offsets = list(range(0, length, stride)) # Prepare Weight cache_key = (segment_length, float(transition_power), mix_dtype) weight = _WEIGHT_CACHE.get(cache_key) if weight is None: weight = mx.concatenate([ mx.arange(1, segment_length // 2 + 1), mx.arange(segment_length - segment_length // 2, 0, -1), ], axis=0) weight = (weight / mx.max(weight)) ** transition_power _WEIGHT_CACHE[cache_key] = weight progress_bar = None if progress_enabled: from tqdm import tqdm progress_bar = tqdm(total=len(offsets), desc="segments", unit="seg", leave=False) # --- BATCHING STATE --- batch_inputs = [] batch_indices = [] if hasattr(model, "valid_length"): std_valid_len = model.valid_length(segment_length) else: std_valid_len = segment_length def flush_batch(): nonlocal batch_inputs, batch_indices, out, sum_weight if not batch_inputs: return # 1. Stack: (Batch_Segments, Audio_Batch, Channels, Time) batch_tensor = mx.stack(batch_inputs) # 2. Reshape: Flatten Segments into Batch -> (Total_Batch, Channels, Time) b_seg, b_audio, channels, length = batch_tensor.shape batch_tensor_flat = batch_tensor.reshape(b_seg * b_audio, channels, length) # 3. Run Model (Standard 3D Input) batch_out_flat = model(batch_tensor_flat) # 4. Unflatten: (Batch_Segments, Audio_Batch, Sources, Channels, Time) _, sources, out_c, out_t = batch_out_flat.shape batch_out = batch_out_flat.reshape(b_seg, b_audio, sources, out_c, out_t) for i, idx in enumerate(batch_indices): chunk_out = center_trim(batch_out[i], segment_length) offset = offsets[idx] end = offset + segment_length out = out.at[:, :, :, offset:end].add(weight.reshape(1, 1, 1, -1) * chunk_out) sum_weight = sum_weight.at[offset:end].add(weight) if progress_bar is not None: progress_bar.update(1) mx.eval(out, sum_weight) batch_inputs = [] batch_indices = [] try: for i, offset in enumerate(offsets): this_chunk_len = min(segment_length, length - offset) chunk = TensorChunk(mix_chunk, offset, this_chunk_len) # Batch only standard-sized chunks if this_chunk_len == segment_length: padded = chunk.padded(std_valid_len) batch_inputs.append(padded) batch_indices.append(i) if len(batch_inputs) >= batch_size: flush_batch() else: # Flush pending batch flush_batch() # Run odd-sized chunk individually if hasattr(model, "valid_length"): valid_len = model.valid_length(this_chunk_len) else: valid_len = this_chunk_len padded = chunk.padded(valid_len) # FIX: Pass 'padded' directly. It is already (Batch, Channels, Time). chunk_out = model(padded) chunk_out = center_trim(chunk_out, this_chunk_len) end = offset + this_chunk_len w = weight[:this_chunk_len].reshape(1, 1, 1, -1) out = out.at[:, :, :, offset:end].add(w * chunk_out) sum_weight = sum_weight.at[offset:end].add(weight[:this_chunk_len]) mx.eval(out, sum_weight) # Eval to bound graph size if progress_bar is not None: progress_bar.update(1) flush_batch() finally: if progress_bar is not None: progress_bar.close() if bool(mx.any(sum_weight == 0).item()): raise ValueError("sum_weight has zeros; check segment and overlap settings") out = out / sum_weight mx.eval(out) return out # No split path valid_length = model.valid_length(length) if hasattr(model, "valid_length") else length padded_mix = mix_chunk.padded(valid_length) out = model(padded_mix) return center_trim(out, length)