""" SonicMaster inference, loaded once per process. The algorithm is the authors' own, from infer_single.py in https://github.com/AMAAI-Lab/SonicMaster (Apache 2.0, vendored here as LICENSE-SonicMaster): the input is resampled to 44.1 kHz stereo, cut into overlapping chunks, encoded by the Stable Audio VAE, carried chunk to chunk through the flow model, then decoded and crossfaded back together. What differs is where the work happens. Their script loads the checkpoint, the text encoder and the VAE on every invocation, which suits a command line and not a Space. Here they load at import, which is also what ZeroGPU wants: a real GPU exists only inside the decorated function, and Hugging Face runs a CUDA emulation outside it so module-level placement still works. """ import os from pathlib import Path import soundfile as sf import torch import torchaudio import yaml from diffusers import AutoencoderOobleck from huggingface_hub import hf_hub_download from safetensors.torch import load_file from model import TangoFlux SAMPLE_RATE = 44100 CHANNELS = 2 WEIGHTS_REPO = "amaai-lab/SonicMaster" WEIGHTS_FILE = "model.safetensors" # The VAE the model decodes through. It is gated, so this Space needs an # HF_TOKEN secret from an account that has accepted the Stability AI Community # License. People calling the Space do not need to accept anything: the # download happens here, with the Space's own credentials. VAE_REPO = "stabilityai/stable-audio-open-1.0" VAE_SUBFOLDER = "vae" HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") CONFIG = Path(__file__).parent / "configs" / "tangoflux_config.yaml" device = "cuda" if torch.cuda.is_available() else "cpu" torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True with CONFIG.open(encoding="utf-8") as stream: _config = yaml.safe_load(stream) model = TangoFlux(config=_config["model"]) model.load_state_dict(load_file(hf_hub_download(WEIGHTS_REPO, WEIGHTS_FILE)), strict=False) model.to(device).eval() for parameter in model.text_encoder.parameters(): parameter.requires_grad = False model.text_encoder.eval() if HF_TOKEN is None: raise RuntimeError( f"{VAE_REPO} is gated, so this Space needs an HF_TOKEN secret from an " f"account that has accepted the Stability AI Community License") vae = AutoencoderOobleck.from_pretrained(VAE_REPO, subfolder=VAE_SUBFOLDER, token=HF_TOKEN) vae.to(device).eval() def _chunks(audio, chunk_size, stride): """ Cut audio into equal chunks, padding the last one. Args: audio (torch.Tensor): Waveform shaped [channels, samples]. chunk_size (int): Samples per chunk. stride (int): Samples between chunk starts. Returns: chunks (list): Tensors shaped [channels, chunk_size]. """ out, start, total = [], 0, audio.shape[1] while start < total: chunk = audio[:, start:min(start + chunk_size, total)] if chunk.shape[1] < chunk_size: chunk = torch.nn.functional.pad(chunk, (0, chunk_size - chunk.shape[1])) out.append(chunk) start += stride return out def _stitch(decoded, overlap): """ Crossfade decoded chunks back into one waveform. Args: decoded (list): Tensors shaped [1, channels, samples]. overlap (int): Samples the neighbors share. Returns: final (torch.Tensor): One waveform shaped [1, channels, samples]. """ final = decoded[0] for chunk in decoded[1:]: fade = torch.linspace(1.0, 0.0, steps=overlap).view(1, 1, -1) blended = final[:, :, -overlap:] * fade + chunk[:, :, :overlap] * (1.0 - fade) final = torch.cat([final[:, :, :-overlap], blended, chunk[:, :, overlap:]], dim=2) return final @torch.no_grad() def enhance(input_path, prompt, output_path, num_inference_steps=10, guidance_scale=1.0, seed=0, solver="Euler", chunk_duration=30, overlap_duration=10, vae_batch_size=10): """ Apply a text-described mastering or restoration edit to one recording. Args: input_path (str): The recording to edit. prompt (str): What to change about it. output_path (str): Where to write the result. num_inference_steps (int): Flow solver steps. guidance_scale (float): How strongly the prompt is followed. seed (int): Seed for the noise, so a run can be repeated. solver (str): "Euler" or "rk4". chunk_duration (int): Seconds the model sees at once. overlap_duration (int): Seconds shared by neighboring chunks, which both condition the next chunk and crossfade the result. vae_batch_size (int): Chunks encoded per VAE batch. Returns: output_path (str): The path written. """ audio, rate = torchaudio.load(str(input_path)) if audio.shape[0] == 1: audio = audio.repeat(CHANNELS, 1) elif audio.shape[0] > CHANNELS: audio = audio[:CHANNELS, :] if rate != SAMPLE_RATE: audio = torchaudio.functional.resample(audio, rate, SAMPLE_RATE) chunk_size = chunk_duration * SAMPLE_RATE overlap = overlap_duration * SAMPLE_RATE if not 0 < overlap < chunk_size: raise ValueError("the overlap must be shorter than the chunk") chunks = _chunks(audio.to(device), chunk_size, chunk_size - overlap) stacked = torch.stack(chunks) latents = [vae.encode(stacked[start:start + vae_batch_size].to(device)) .latent_dist.mode() for start in range(0, stacked.shape[0], vae_batch_size)] degraded = torch.cat(latents, dim=0) decoded, carried = [], None for index in range(degraded.shape[0]): result = model.inference_flow( degraded[index].unsqueeze(0).transpose(1, 2), prompt, audiocond_latents=carried, num_inference_steps=num_inference_steps, timesteps=None, guidance_scale=guidance_scale, duration=chunk_duration, seed=seed, disable_progress=True, num_samples_per_prompt=1, callback_on_step_end=None, solver=solver, ) waveform = torch.clamp(vae.decode(result.transpose(2, 1)).sample.cpu(), -1.0, 1.0) decoded.append(waveform) carried = vae.encode(waveform[:, :, -overlap:].to(device)) \ .latent_dist.mode().transpose(1, 2) final = _stitch(decoded, overlap) sf.write(str(output_path), final.squeeze(0).numpy().T, SAMPLE_RATE, format="WAV") return str(output_path)