sonicmaster_runtime.py

"""
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)