Spaces:
Running on Zero
Running on Zero
| """ | |
| Audio utility functions shared between test and validation. | |
| """ | |
| import numpy as np | |
| import torch | |
| def normalize_audio(audio_np: np.ndarray) -> np.ndarray: | |
| """Normalize audio to [-0.9, 0.9] range.""" | |
| max_val = np.abs(audio_np).max() | |
| if max_val > 0: | |
| return audio_np / max_val * 0.9 | |
| return audio_np | |
| def decode_reference_audio(sample: dict, audio_vae, device: str, sample_rate: int = 44100) -> np.ndarray | None: | |
| """ | |
| Decode reference audio from a preprocessed sample. | |
| Preprocessed data stores audio as VAE latent features in 'audio_feats' column. | |
| We use the 'audio_mask' to extract only valid audio frames (skipping text-position padding). | |
| Args: | |
| sample: A single dataset sample dict containing 'audio_feats' and optionally 'audio_mask'. | |
| audio_vae: The AudioVAE model instance. | |
| device: Device string (e.g. "cuda"). | |
| sample_rate: Audio sample rate (unused in decoding, kept for API consistency). | |
| Returns: | |
| 1-D numpy array of decoded audio, or None on failure. | |
| """ | |
| from einops import rearrange | |
| if "audio_feats" not in sample: | |
| return None | |
| audio_feats = sample["audio_feats"] | |
| audio_mask = sample.get("audio_mask", None) | |
| # Convert to tensor if needed (Arrow datasets return lists) | |
| if isinstance(audio_feats, list): | |
| audio_feats = torch.tensor(audio_feats, dtype=torch.float32) | |
| if audio_mask is not None and isinstance(audio_mask, list): | |
| audio_mask = torch.tensor(audio_mask, dtype=torch.int32) | |
| if not isinstance(audio_feats, torch.Tensor): | |
| return None | |
| # Extract only valid audio frames using audio_mask | |
| if audio_mask is not None and isinstance(audio_mask, torch.Tensor): | |
| audio_indices = (audio_mask == 1).nonzero(as_tuple=True)[0] | |
| if len(audio_indices) > 0: | |
| audio_feats = audio_feats[audio_indices] | |
| else: | |
| return None | |
| # Reverse preprocessing: [T, P, D] -> [1, T*P, D] -> [1, D, T*P] -> decode | |
| audio_feats = rearrange(audio_feats, "t p d -> 1 (t p) d") | |
| audio_feats = audio_feats.transpose(1, 2) # [1, D, T*P] | |
| audio_feats = audio_feats.to(device).float() | |
| with torch.no_grad(): | |
| ref_audio = audio_vae.decode(audio_feats) # [1, 1, samples] | |
| return ref_audio.cpu().float().numpy().flatten() | |
| def decode_latents_batched( | |
| latents, | |
| audio_vae, | |
| device: str, | |
| sample_rate: int = 44100, | |
| max_batch_tokens: int = 0, | |
| ) -> list: | |
| """Decode a list of ``(T, P, D)`` VAE latents to 1-D fp32 mono waveforms in | |
| batches. | |
| Variable-length latents are right-padded to the per-batch max along the | |
| ``(T*P)`` time axis, decoded as one ``[B, D, L_max]`` tensor, then each | |
| decoded waveform is trimmed back to its true length ``L_i * hop_length``. | |
| The AudioVAE decoder is a *causal* conv stack (output sample ``t`` depends | |
| only on inputs ``<= t``), so right-padding + trimming is numerically | |
| identical to a ``B=1`` decode up to fp reduction-order jitter (~1e-6). | |
| Args: | |
| latents: list of ``(T, P, D)`` tensors (bf16/fp32) or lists; ``None`` | |
| entries pass through as ``None``. | |
| audio_vae: AudioVAE (v1) or AudioVAEV2 instance. | |
| device: device string for the decode. | |
| sample_rate: fallback output sample rate (only used if the VAE does not | |
| expose ``out_sample_rate``). | |
| max_batch_tokens: in-batch ``T*P`` frame budget; ``<=0`` packs all | |
| latents into a single batch. | |
| Returns: | |
| list aligned to ``latents``; each entry is a 1-D fp32 numpy waveform or | |
| ``None`` (for ``None`` / invalid inputs). | |
| """ | |
| from einops import rearrange | |
| results = [None] * len(latents) | |
| # Flatten valid latents to [D, L] and record original index + length. | |
| items = [] # (orig_idx, z[D, L], L) | |
| for i, lat in enumerate(latents): | |
| if lat is None: | |
| continue | |
| t = lat | |
| if isinstance(t, list): | |
| t = torch.tensor(t, dtype=torch.float32) | |
| if not isinstance(t, torch.Tensor): | |
| continue | |
| z = rearrange(t, "tt p d -> (tt p) d").transpose(0, 1).contiguous() # [D, T*P] | |
| items.append((i, z, int(z.shape[-1]))) | |
| if not items: | |
| return results | |
| is_v2 = getattr(audio_vae, "sr_bin_boundaries", None) is not None | |
| out_sr = int(getattr(audio_vae, "out_sample_rate", sample_rate)) | |
| # Sort by length so each padded batch packs similar-length latents tightly. | |
| items.sort(key=lambda x: x[2]) | |
| # Group into batches by *padded* volume. Because items are length-sorted | |
| # ascending, the newest item's length is the batch max, so | |
| # ``L * (len(cur)+1)`` is exactly the padded ``[B, D, L_max]`` token count — | |
| # budgeting on that (not the sum of true lengths) bounds the compute wasted | |
| # on zero-padding when a batch mixes short and long latents. | |
| batches, cur = [], [] | |
| for it in items: | |
| L = it[2] | |
| if cur and max_batch_tokens > 0 and L * (len(cur) + 1) > max_batch_tokens: | |
| batches.append(cur) | |
| cur = [] | |
| cur.append(it) | |
| if cur: | |
| batches.append(cur) | |
| for batch in batches: | |
| L_max = max(it[2] for it in batch) | |
| B = len(batch) | |
| D = batch[0][1].shape[0] | |
| z = torch.zeros(B, D, L_max, dtype=torch.float32) | |
| for bi, (_, zz, L) in enumerate(batch): | |
| z[bi, :, :L] = zz.float() | |
| z = z.to(device) | |
| with torch.no_grad(): | |
| if is_v2: | |
| sr_cond = torch.full((B,), out_sr, device=device, dtype=torch.int32) | |
| out = audio_vae.decode(z, sr_cond) # [B, 1, S_max] | |
| else: | |
| out = audio_vae.decode(z) # [B, 1, S_max] | |
| out = out.cpu().float() | |
| S_max = out.shape[-1] | |
| # The decoder upsamples every latent frame by a fixed integer factor, so | |
| # the per-frame sample count is exactly S_max // L_max (architecture- and | |
| # sample-rate-agnostic; works for v1 where in_sr==out_sr and for v2 where | |
| # the output is resampled to out_sample_rate). Trimming each row to its | |
| # true length is what keeps batching lossless vs the B=1 decode. | |
| ratio = S_max // L_max | |
| for bi, (orig_idx, _, L) in enumerate(batch): | |
| n_samp = min(L * ratio, S_max) | |
| results[orig_idx] = out[bi, 0, :n_samp].numpy().flatten() | |
| return results | |