Spaces:

teamup-tech
/

FxNorm-Automix

Running on Zero

Add domain tags for the Home tab filters of HARP

4bbeb75 verified 6 days ago

17.5 kB

import spaces  # imported first, as ZeroGPU requires

import sys
sys.stdout.reconfigure(line_buffering=True)
sys.setrecursionlimit(int(1e6))

import functools
import tempfile
import numpy as np
import scipy.signal
import torch
import soundfile as sf
import librosa
import gradio as gr

from pymixconsole.parameter import Parameter
from pymixconsole.parameter_list import ParameterList
from pyharp import ModelCard, build_endpoint
from pyharp.tags import Channels, SampleRate, Subcategory

from automix.common_dataprocessing import load_wav, save_wav, generate_data
from config import *
from automix.common_supernet import SuperNet
from automix.common_datatypes import DataType
from automix.utils_data_normalization import (
    get_comp_matching, get_eq_matching, lufs_normalize,
    get_panning_matching, amp_to_db, get_reverb_send,
)

model_card = ModelCard(
    name="FxNorm-Automix",
    description=(
        "Automatic music mixing with deep learning and out-of-domain data. "
        "Accepts up to 4 stems (vocals, bass, drums, other), each up to 3.5 "
        "minutes. Vocals are required, and bass, drums, and other are optional. The stems "
        "are mixed at 44.1 kHz in stereo, and the mix comes back at the highest rate among them. "
        "Repository: https://github.com/sony/FxNorm-automix Paper: "
        "https://arxiv.org/abs/2208.11428"
    ),
    author="Marco A. Martínez-Ramírez, Wei-Hsiang Liao, Giorgio Fabbro, Stefan Uhlich, Chihiro Nagashima, Yuki Mitsufuji",
    tags=[Subcategory.AUTOMATIC_MIXING, SampleRate(SR), Channels(N_CHANNELS), 'music', 'domain:music'],
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
USE_AMP = torch.cuda.is_available()

EQ_PARAMETERS = ParameterList()
EQ_PARAMETERS.add(Parameter('low_shelf_gain', EQ_GAIN, 'float', minimum=EQ_GAIN, maximum=EQ_GAIN))
EQ_PARAMETERS.add(Parameter('low_shelf_freq', 600.0, 'float', minimum=500.0, maximum=700.0))
EQ_PARAMETERS.add(Parameter('high_shelf_gain', EQ_GAIN, 'float', minimum=EQ_GAIN, maximum=EQ_GAIN))
EQ_PARAMETERS.add(Parameter('high_shelf_freq', 8500.0, 'float', minimum=7000.0, maximum=10000.0))

PRE_REVERB_PARAMETERS = ParameterList()
PRE_REVERB_PARAMETERS.add(Parameter('wet', 0.5, 'float', minimum=0.5, maximum=0.5))
PRE_REVERB_PARAMETERS.add(Parameter('dry', 0.5, 'float', minimum=0.5, maximum=0.5))
PRE_REVERB_PARAMETERS.add(Parameter('decay', 1.0, 'float', minimum=1.0, maximum=1.0))
PRE_REVERB_PARAMETERS.add(Parameter('pre_delay', 0, 'int', units='ms', minimum=0, maximum=0))

REVERB_PARAMETERS = ParameterList()
REVERB_PARAMETERS.add(Parameter('wet', 1.0, 'float', minimum=1.0, maximum=1.0))
REVERB_PARAMETERS.add(Parameter('dry', 0.0, 'float', minimum=0.0, maximum=0.0))
REVERB_PARAMETERS.add(Parameter('decay', 1.0, 'float', minimum=1.0, maximum=1.0))
REVERB_PARAMETERS.add(Parameter('pre_delay', 0, 'int', units='ms', minimum=0, maximum=0))


def _load_ir(path):
    data = load_wav(path)
    return {'impulse_response': functools.partial(generate_data, file_path_or_data=data)}


# IR and pre-IR are used from OpenAIR library
# Abies Grandis Forest, Wheldrake Wood (s1r1 and s1r2)
# Lady Chapel, St Albans Cathedral (stalbans_a_ortf, stalbans_b_ortf)
# Stairway, University of York (stairwell_ortf)
PRE_IR = [_load_ir('pre-ir/S1R1_SF.wav'), _load_ir('pre-ir/S1R2_SF.wav')]
IR = [_load_ir('ir/stalbans_a_ortf.wav'), _load_ir('ir/stalbans_b_ortf.wav'), _load_ir('ir/stairwell_ortf.wav')]

PRE_REVERB_PARAMETERS.add(Parameter('index', 0, 'int', minimum=0, maximum=len(PRE_IR)))
REVERB_PARAMETERS.add(Parameter('index', 0, 'int', minimum=0, maximum=len(IR)))

# --- Load model ---
# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the model is built there alone and placed on DEVICE on its main thread, where
# ZeroGPU's emulation records the placement, rather than inside a GPU call
if __name__ != "__main__":
    net = torch.load('checkpoints/net_mixture.dump', map_location='cpu', weights_only=False)
    net.load_state_dict(torch.load('checkpoints/current_model_for_mixture.params',
                                   map_location='cpu', weights_only=False))

    super_net = SuperNet(
        net,
        stft_window=torch.from_numpy(STFT_WINDOW.astype(np.float32)),
        stft_hop_length=HOP_LENGTH_CONFIG,
        batched_valid=False,
        unfolding_params=None,
        training_length=TRAINING_SEQ_LENGTH,
        training_batch_size=BATCH_SIZE,
        use_amp=USE_AMP,
    )
    super_net.eval()
    super_net.to(DEVICE)

# --- Load and smooth features ---
features_mean = np.load('checkpoints/features_MUSDB18.npy', allow_pickle=True)[()]

for _effect in EFFECTS:
    if _effect == 'eq':
        for _key in STEMS:
            window_length = 401 if _key in ['other', 'vocals'] else 151
            features_mean[_effect][_key] = scipy.signal.savgol_filter(
                features_mean[_effect][_key], window_length, 1, mode='mirror')
    elif _effect == 'panning':
        for _key in STEMS:
            features_mean[_effect][_key] = scipy.signal.savgol_filter(
                features_mean[_effect][_key], 501, 1, mode='mirror')


def normalize_audio_wave(stem_args):
    """
    Effect Normalizes audio input
    Args:
        tuple: audio x list of effects to normalize x source (e.g., vocals) x dataset features
    Returns:
        normalized audio
    """

    audio = stem_args[0]
    effects = stem_args[1]
    src = stem_args[2]
    features_mean_ = stem_args[3]

    print(f'{src} normalizing {effects}...')
    all_zeros = not np.any(audio)

    if not all_zeros:
        audio_track = np.pad(audio, ((FFT_SIZE, FFT_SIZE), (0, 0)))
        assert len(audio_track.shape) == 2
        if audio_track.shape[1] == 1:
            audio_track = np.repeat(audio_track, 2, axis=-1)

        output_audio = audio_track.copy()
        max_db = amp_to_db(np.max(np.abs(output_audio)))

        if max_db > MIN_DB:
            for effect in effects:
                if effect == 'eq':
                    for ch in range(audio_track.shape[1]):
                        audio_eq_matched = get_eq_matching(
                                            output_audio[:, ch], 
                                            features_mean_[effect][src],
                                            sr=SR, n_fft=FFT_SIZE, 
                                            hop_length=HOP_LENGTH,
                                            min_db=MIN_DB, 
                                            ntaps=NTAPS, 
                                            lufs=LUFS)
                        np.copyto(output_audio[:, ch], audio_eq_matched)

                elif effect == 'compression':
                    assert len(features_mean_[effect][src]) == 2
                    for ch in range(audio_track.shape[1]):
                        audio_comp_matched = get_comp_matching(
                            output_audio[:, ch],
                            features_mean_[effect][src][0],
                            features_mean_[effect][src][1],
                            comp_settings[src]['ratio'],
                            comp_settings[src]['attack'],
                            comp_settings[src]['release'],
                            sr=SR, min_db=MIN_DB, min_th=COMP_MIN_TH,
                            comp_peak_norm=COMP_PEAK_NORM, max_ratio=COMP_MAX_RATIO,
                            n_mels=comp_settings[src]['n_mels'],
                            true_peak=COMP_TRUE_PEAK, percentile=COMP_PERCENTILE,
                            expander=COMP_USE_EXPANDER)
                        np.copyto(output_audio[:, ch], audio_comp_matched[:, 0])

                elif effect == 'panning':
                    if PANNING_FRAMES:
                        ref_phi = features_mean_[effect][src][:-1]
                        bins_per_frame = int(2 * (features_mean_[effect][src].shape[0] - 1) / FFT_SIZE_PAN)
                        ref_phi = ref_phi.reshape(-1, bins_per_frame).mean(axis=1)
                        ref_phi = np.append(ref_phi, features_mean_[effect][src][-1])
                    else:
                        ref_phi = features_mean_[effect][src]

                    MAX_FREQ_PAN = 6000 if src == 'drums' else 16000

                    audio_pan_matched = get_panning_matching(output_audio, 
                                                             PANNING_GAIN_K * ref_phi,
                                                             sr=SR, 
                                                             n_fft=FFT_SIZE_PAN, 
                                                             hop_length=HOP_LENGTH_PAN,
                                                             min_db_f=MIN_DB_f, 
                                                             max_freq_pan=MAX_FREQ_PAN, 
                                                             frames=PANNING_FRAMES)
                    np.copyto(output_audio, audio_pan_matched)

                elif effect == 'loudness':
                    output_audio = lufs_normalize(output_audio, SR, features_mean_[effect][src])

                elif effect == 'reverb':
                    if src not in ['bass', 'drums']:
                        audio_reverb_send = get_reverb_send(
                            output_audio, EQ_PARAMETERS, REVERB_PARAMETERS, IR,
                            bands=EQ_BANDS_VOCALS_OTHER,
                            eq_prob=EQ_PROBABILITY, rv_prob=REVERB_PROBABILITY,
                            parallel=FX_PARALLEL, shuffle=FX_SHUFFLE, sr=SR)
                        np.copyto(output_audio, audio_reverb_send)

                elif effect == 'prereverb':
                    if src not in ['bass', 'drums']:
                        audio_reverb_send = get_reverb_send(
                            output_audio, EQ_PARAMETERS, PRE_REVERB_PARAMETERS, PRE_IR,
                            bands=EQ_BANDS_VOCALS_OTHER,
                            eq_prob=PRE_EQ_PROBABILITY, rv_prob=PRE_REVERB_PROBABILITY,
                            parallel=PRE_FX_PARALLEL, shuffle=PRE_FX_SHUFFLE, sr=SR)
                        np.copyto(output_audio, audio_reverb_send)

        output_audio = output_audio[FFT_SIZE:FFT_SIZE + audio.shape[0]]
    else:
        print(f'{src} is only zeros...')
        output_audio = audio

    return output_audio


def process_fn(vocals_path, bass_path, drums_path, other_path):
    # Checked before the GPU is reserved, so stems the model cannot take spend none of the
    # caller's ZeroGPU allowance
    paths = {'vocals': vocals_path, 'bass': bass_path, 'drums': drums_path, 'other': other_path}

    if not any(paths.values()):
        raise gr.Error("Please provide at least one stem.")

    for stem in STEMS:
        if paths[stem]:
            try:
                info = sf.info(paths[stem])
            except Exception as exc:
                raise gr.Error(f"{stem}: could not read the audio file: {exc}") from exc
            if info.frames == 0:
                raise gr.Error(f"{stem}: the audio file is empty.")
            if info.channels not in (1, N_CHANNELS):
                raise gr.Error(f'{stem}: expected {N_CHANNELS} channels, got {info.channels}')
            if info.duration > MAX_STEM_SECONDS:
                raise gr.Error(
                    f"{stem}: the stem is {info.duration:.1f} s long. Trim it to "
                    f"{MAX_STEM_SECONDS / 60:g} minutes or less."
                )

    return mix(paths)


# Every stem is padded with 30 s of context, and the model takes at most MAX_LENGTH seconds with
# the padding, so a longer stem would come back cut short
CONTEXT_SECONDS = 30
MAX_STEM_SECONDS = MAX_LENGTH - CONTEXT_SECONDS


def gpu_seconds(paths):
    # Measured at about 29 s for four 30 s stems and for one 60 s stem, padded
    longest = max(sf.info(path).duration for path in paths.values() if path)
    return 15 + 0.4 * (longest + CONTEXT_SECONDS)


@spaces.GPU(duration=gpu_seconds)
def mix(paths):
    # --- Load and validate stems ---
    pad_samples = CONTEXT_SECONDS * SR
    audio = {}
    samples = []
    subtype_list = []
    samplingrate_list = []

    for stem in STEMS:
        if paths[stem]:
            # soundfile reads every format HARP sends, where the authors' reader takes only
            # integer PCM WAV, and returns the integer PCM the pipeline below expects
            _arr, _sr = sf.read(paths[stem], dtype="int32", always_2d=True)
            samplingrate_list.append(_sr)
            # A lossy or floating-point input has no PCM depth for the WAV output to keep
            subtype = sf.info(paths[stem]).subtype
            subtype_list.append(subtype if subtype.startswith("PCM_") else "PCM_24")

            # ensure stereo
            if _arr.shape[1] == 1:
                _arr = np.repeat(_arr, 2, axis=-1)
            elif _arr.shape[1] != N_CHANNELS:
                raise gr.Error(f'{stem}: expected {N_CHANNELS} channels, got {_arr.shape[1]}')

            if _sr != SR:
                # resample to model's expected SR (44100); librosa needs float32 in [-1, 1]
                # so we convert int16 to float, resample, then convert back to int16
                temp = _arr.astype(np.float32) / (1. + np.iinfo(_arr.dtype).max)
                temp = librosa.resample(temp.T, orig_sr=_sr, target_sr=SR,
                                        res_type='kaiser_best').T
                temp = temp * (1 + np.iinfo(np.int16).max)
                _arr = np.clip(temp, np.iinfo(np.int16).min,
                               np.iinfo(np.int16).max).astype(np.int16)

            # Pad the start with 30s of wrapped audio so the model has context at
            # t=0. This follows the resample rather than preceding it, because
            # pad_samples counts samples at SR: padding first adds that many
            # samples at the input's own rate, and resampling then leaves a pad
            # shorter than the trim after inference removes. A 48 kHz input came
            # back empty for exactly that reason.
            _arr = np.pad(_arr, [(pad_samples, 0), (0, 0)], 'wrap')

            samples.append(_arr.shape[0])
            audio[stem] = (SR, _arr)
        else:
            # missing stems are silent so the model always receives 4 tracks
            audio[stem] = (SR, np.zeros((SR, N_CHANNELS), dtype=np.int16))

    max_samples = min(max(samples), int(MAX_LENGTH * SR))
    max_samplingrate = max(samplingrate_list)
    max_subtype = max(subtype_list)

    # --- Build input tensor ---
    data = np.zeros((len(STEMS), 1, max_samples, N_CHANNELS), dtype=np.float32)
    for k, stem in enumerate(STEMS):
        stem_audio = generate_data(audio[stem])
        if stem_audio[:max_samples].shape[0] < max_samples:
            zeros_ = np.zeros_like(data[k][0])
            zeros_[:stem_audio[:max_samples].shape[0]] = stem_audio[:max_samples]
            data[k] = zeros_[:max_samples]
        else:
            data[k] = stem_audio[:max_samples]

    # round up to the next multiple of KERNEL_SIZE_ENCODER as required by the encoder
    new_samples = (1 + max_samples // KERNEL_SIZE_ENCODER) * KERNEL_SIZE_ENCODER
    data = np.pad(data, [(0, 0), (0, 0), (0, new_samples - max_samples), (0, 0)])

    # --- Effect normalization ---
    # changed from org. inference #512 because parallelization may not be safe on HF's CPU tier
    for k, stem in enumerate(STEMS):
        data[k][0] = normalize_audio_wave((data[k][0], EFFECTS, stem, features_mean))

    # --- Run model inference ---
    test_data = torch.from_numpy(data).to(DEVICE)
    with torch.no_grad():
        test_out = super_net.inference(test_data)
        audio_out = test_out[DataType.TIME_SAMPLES].cpu().numpy()
        audio_out = audio_out[..., :max_samples]  # trim encoder padding
        audio_out = audio_out[..., pad_samples:]  # trim the 30s context pad added at load

    # first output stem is the mix
    audio_output = audio_out[0, 0, ...]

    # --- Resample and save output ---
    if max_samplingrate != SR:
        audio_output = librosa.resample(audio_output, orig_sr=SR,
                                        target_sr=max_samplingrate,
                                        res_type='kaiser_best')

    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f:
        out_path = f.name
    save_wav(out_path, max_samplingrate, audio_output.T, subtype=max_subtype)

    return out_path


if __name__ == "__main__":
    with gr.Blocks(title="FxNorm Automix") as demo:
        input_components = [
            gr.Audio(type="filepath", label="Vocals").harp_required(True).set_info("The vocal stem, up to 3.5 minutes. Required."),
            gr.Audio(type="filepath", label="Bass").harp_required(False).set_info((
                "The bass stem. Optional, and left out if the song has none."
            )),
            gr.Audio(type="filepath", label="Drums").harp_required(False).set_info("The drum stem. Optional."),
            gr.Audio(type="filepath", label="Other").harp_required(False).set_info((
                "Everything else: guitars, keys, strings. Optional."
            )),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Mix").set_info("Automatic mix of the input stems."),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )

    demo.queue().launch(show_error=True, pwa=True)

	import spaces # imported first, as ZeroGPU requires

	import sys
	sys.stdout.reconfigure(line_buffering=True)
	sys.setrecursionlimit(int(1e6))

	import functools
	import tempfile
	import numpy as np
	import scipy.signal
	import torch
	import soundfile as sf
	import librosa
	import gradio as gr

	from pymixconsole.parameter import Parameter
	from pymixconsole.parameter_list import ParameterList
	from pyharp import ModelCard, build_endpoint
	from pyharp.tags import Channels, SampleRate, Subcategory

	from automix.common_dataprocessing import load_wav, save_wav, generate_data
	from config import *
	from automix.common_supernet import SuperNet
	from automix.common_datatypes import DataType
	from automix.utils_data_normalization import (
	get_comp_matching, get_eq_matching, lufs_normalize,
	get_panning_matching, amp_to_db, get_reverb_send,
	)

	model_card = ModelCard(
	name="FxNorm-Automix",
	description=(
	"Automatic music mixing with deep learning and out-of-domain data. "
	"Accepts up to 4 stems (vocals, bass, drums, other), each up to 3.5 "
	"minutes. Vocals are required, and bass, drums, and other are optional. The stems "
	"are mixed at 44.1 kHz in stereo, and the mix comes back at the highest rate among them. "
	"Repository: https://github.com/sony/FxNorm-automix Paper: "
	"https://arxiv.org/abs/2208.11428"
	),
	author="Marco A. Martínez-Ramírez, Wei-Hsiang Liao, Giorgio Fabbro, Stefan Uhlich, Chihiro Nagashima, Yuki Mitsufuji",
	tags=[Subcategory.AUTOMATIC_MIXING, SampleRate(SR), Channels(N_CHANNELS), 'music', 'domain:music'],
	)

	DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
	USE_AMP = torch.cuda.is_available()

	EQ_PARAMETERS = ParameterList()
	EQ_PARAMETERS.add(Parameter('low_shelf_gain', EQ_GAIN, 'float', minimum=EQ_GAIN, maximum=EQ_GAIN))
	EQ_PARAMETERS.add(Parameter('low_shelf_freq', 600.0, 'float', minimum=500.0, maximum=700.0))
	EQ_PARAMETERS.add(Parameter('high_shelf_gain', EQ_GAIN, 'float', minimum=EQ_GAIN, maximum=EQ_GAIN))
	EQ_PARAMETERS.add(Parameter('high_shelf_freq', 8500.0, 'float', minimum=7000.0, maximum=10000.0))

	PRE_REVERB_PARAMETERS = ParameterList()
	PRE_REVERB_PARAMETERS.add(Parameter('wet', 0.5, 'float', minimum=0.5, maximum=0.5))
	PRE_REVERB_PARAMETERS.add(Parameter('dry', 0.5, 'float', minimum=0.5, maximum=0.5))
	PRE_REVERB_PARAMETERS.add(Parameter('decay', 1.0, 'float', minimum=1.0, maximum=1.0))
	PRE_REVERB_PARAMETERS.add(Parameter('pre_delay', 0, 'int', units='ms', minimum=0, maximum=0))

	REVERB_PARAMETERS = ParameterList()
	REVERB_PARAMETERS.add(Parameter('wet', 1.0, 'float', minimum=1.0, maximum=1.0))
	REVERB_PARAMETERS.add(Parameter('dry', 0.0, 'float', minimum=0.0, maximum=0.0))
	REVERB_PARAMETERS.add(Parameter('decay', 1.0, 'float', minimum=1.0, maximum=1.0))
	REVERB_PARAMETERS.add(Parameter('pre_delay', 0, 'int', units='ms', minimum=0, maximum=0))


	def _load_ir(path):
	data = load_wav(path)
	return {'impulse_response': functools.partial(generate_data, file_path_or_data=data)}


	# IR and pre-IR are used from OpenAIR library
	# Abies Grandis Forest, Wheldrake Wood (s1r1 and s1r2)
	# Lady Chapel, St Albans Cathedral (stalbans_a_ortf, stalbans_b_ortf)
	# Stairway, University of York (stairwell_ortf)
	PRE_IR = [_load_ir('pre-ir/S1R1_SF.wav'), _load_ir('pre-ir/S1R2_SF.wav')]
	IR = [_load_ir('ir/stalbans_a_ortf.wav'), _load_ir('ir/stalbans_b_ortf.wav'), _load_ir('ir/stairwell_ortf.wav')]

	PRE_REVERB_PARAMETERS.add(Parameter('index', 0, 'int', minimum=0, maximum=len(PRE_IR)))
	REVERB_PARAMETERS.add(Parameter('index', 0, 'int', minimum=0, maximum=len(IR)))

	# --- Load model ---
	# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
	# server, so the model is built there alone and placed on DEVICE on its main thread, where
	# ZeroGPU's emulation records the placement, rather than inside a GPU call
	if __name__ != "__main__":
	net = torch.load('checkpoints/net_mixture.dump', map_location='cpu', weights_only=False)
	net.load_state_dict(torch.load('checkpoints/current_model_for_mixture.params',
	map_location='cpu', weights_only=False))

	super_net = SuperNet(
	net,
	stft_window=torch.from_numpy(STFT_WINDOW.astype(np.float32)),
	stft_hop_length=HOP_LENGTH_CONFIG,
	batched_valid=False,
	unfolding_params=None,
	training_length=TRAINING_SEQ_LENGTH,
	training_batch_size=BATCH_SIZE,
	use_amp=USE_AMP,
	)
	super_net.eval()
	super_net.to(DEVICE)

	# --- Load and smooth features ---
	features_mean = np.load('checkpoints/features_MUSDB18.npy', allow_pickle=True)[()]

	for _effect in EFFECTS:
	if _effect == 'eq':
	for _key in STEMS:
	window_length = 401 if _key in ['other', 'vocals'] else 151
	features_mean[_effect][_key] = scipy.signal.savgol_filter(
	features_mean[_effect][_key], window_length, 1, mode='mirror')
	elif _effect == 'panning':
	for _key in STEMS:
	features_mean[_effect][_key] = scipy.signal.savgol_filter(
	features_mean[_effect][_key], 501, 1, mode='mirror')


	def normalize_audio_wave(stem_args):
	"""
	Effect Normalizes audio input

	Args:
	tuple: audio x list of effects to normalize x source (e.g., vocals) x dataset features

	Returns:
	normalized audio
	"""

	audio = stem_args[0]
	effects = stem_args[1]
	src = stem_args[2]
	features_mean_ = stem_args[3]

	print(f'{src} normalizing {effects}...')
	all_zeros = not np.any(audio)

	if not all_zeros:
	audio_track = np.pad(audio, ((FFT_SIZE, FFT_SIZE), (0, 0)))
	assert len(audio_track.shape) == 2
	if audio_track.shape[1] == 1:
	audio_track = np.repeat(audio_track, 2, axis=-1)

	output_audio = audio_track.copy()
	max_db = amp_to_db(np.max(np.abs(output_audio)))

	if max_db > MIN_DB:
	for effect in effects:
	if effect == 'eq':
	for ch in range(audio_track.shape[1]):
	audio_eq_matched = get_eq_matching(
	output_audio[:, ch],
	features_mean_[effect][src],
	sr=SR, n_fft=FFT_SIZE,
	hop_length=HOP_LENGTH,
	min_db=MIN_DB,
	ntaps=NTAPS,
	lufs=LUFS)
	np.copyto(output_audio[:, ch], audio_eq_matched)

	elif effect == 'compression':
	assert len(features_mean_[effect][src]) == 2
	for ch in range(audio_track.shape[1]):
	audio_comp_matched = get_comp_matching(
	output_audio[:, ch],
	features_mean_[effect][src][0],
	features_mean_[effect][src][1],
	comp_settings[src]['ratio'],
	comp_settings[src]['attack'],
	comp_settings[src]['release'],
	sr=SR, min_db=MIN_DB, min_th=COMP_MIN_TH,
	comp_peak_norm=COMP_PEAK_NORM, max_ratio=COMP_MAX_RATIO,
	n_mels=comp_settings[src]['n_mels'],
	true_peak=COMP_TRUE_PEAK, percentile=COMP_PERCENTILE,
	expander=COMP_USE_EXPANDER)
	np.copyto(output_audio[:, ch], audio_comp_matched[:, 0])

	elif effect == 'panning':
	if PANNING_FRAMES:
	ref_phi = features_mean_[effect][src][:-1]
	bins_per_frame = int(2 * (features_mean_[effect][src].shape[0] - 1) / FFT_SIZE_PAN)
	ref_phi = ref_phi.reshape(-1, bins_per_frame).mean(axis=1)
	ref_phi = np.append(ref_phi, features_mean_[effect][src][-1])
	else:
	ref_phi = features_mean_[effect][src]

	MAX_FREQ_PAN = 6000 if src == 'drums' else 16000

	audio_pan_matched = get_panning_matching(output_audio,
	PANNING_GAIN_K * ref_phi,
	sr=SR,
	n_fft=FFT_SIZE_PAN,
	hop_length=HOP_LENGTH_PAN,
	min_db_f=MIN_DB_f,
	max_freq_pan=MAX_FREQ_PAN,
	frames=PANNING_FRAMES)
	np.copyto(output_audio, audio_pan_matched)

	elif effect == 'loudness':
	output_audio = lufs_normalize(output_audio, SR, features_mean_[effect][src])

	elif effect == 'reverb':
	if src not in ['bass', 'drums']:
	audio_reverb_send = get_reverb_send(
	output_audio, EQ_PARAMETERS, REVERB_PARAMETERS, IR,
	bands=EQ_BANDS_VOCALS_OTHER,
	eq_prob=EQ_PROBABILITY, rv_prob=REVERB_PROBABILITY,
	parallel=FX_PARALLEL, shuffle=FX_SHUFFLE, sr=SR)
	np.copyto(output_audio, audio_reverb_send)

	elif effect == 'prereverb':
	if src not in ['bass', 'drums']:
	audio_reverb_send = get_reverb_send(
	output_audio, EQ_PARAMETERS, PRE_REVERB_PARAMETERS, PRE_IR,
	bands=EQ_BANDS_VOCALS_OTHER,
	eq_prob=PRE_EQ_PROBABILITY, rv_prob=PRE_REVERB_PROBABILITY,
	parallel=PRE_FX_PARALLEL, shuffle=PRE_FX_SHUFFLE, sr=SR)
	np.copyto(output_audio, audio_reverb_send)

	output_audio = output_audio[FFT_SIZE:FFT_SIZE + audio.shape[0]]
	else:
	print(f'{src} is only zeros...')
	output_audio = audio

	return output_audio


	def process_fn(vocals_path, bass_path, drums_path, other_path):
	# Checked before the GPU is reserved, so stems the model cannot take spend none of the
	# caller's ZeroGPU allowance
	paths = {'vocals': vocals_path, 'bass': bass_path, 'drums': drums_path, 'other': other_path}

	if not any(paths.values()):
	raise gr.Error("Please provide at least one stem.")

	for stem in STEMS:
	if paths[stem]:
	try:
	info = sf.info(paths[stem])
	except Exception as exc:
	raise gr.Error(f"{stem}: could not read the audio file: {exc}") from exc
	if info.frames == 0:
	raise gr.Error(f"{stem}: the audio file is empty.")
	if info.channels not in (1, N_CHANNELS):
	raise gr.Error(f'{stem}: expected {N_CHANNELS} channels, got {info.channels}')
	if info.duration > MAX_STEM_SECONDS:
	raise gr.Error(
	f"{stem}: the stem is {info.duration:.1f} s long. Trim it to "
	f"{MAX_STEM_SECONDS / 60:g} minutes or less."
	)

	return mix(paths)


	# Every stem is padded with 30 s of context, and the model takes at most MAX_LENGTH seconds with
	# the padding, so a longer stem would come back cut short
	CONTEXT_SECONDS = 30
	MAX_STEM_SECONDS = MAX_LENGTH - CONTEXT_SECONDS


	def gpu_seconds(paths):
	# Measured at about 29 s for four 30 s stems and for one 60 s stem, padded
	longest = max(sf.info(path).duration for path in paths.values() if path)
	return 15 + 0.4 * (longest + CONTEXT_SECONDS)


	@spaces.GPU(duration=gpu_seconds)
	def mix(paths):
	# --- Load and validate stems ---
	pad_samples = CONTEXT_SECONDS * SR
	audio = {}
	samples = []
	subtype_list = []
	samplingrate_list = []

	for stem in STEMS:
	if paths[stem]:
	# soundfile reads every format HARP sends, where the authors' reader takes only
	# integer PCM WAV, and returns the integer PCM the pipeline below expects
	_arr, _sr = sf.read(paths[stem], dtype="int32", always_2d=True)
	samplingrate_list.append(_sr)
	# A lossy or floating-point input has no PCM depth for the WAV output to keep
	subtype = sf.info(paths[stem]).subtype
	subtype_list.append(subtype if subtype.startswith("PCM_") else "PCM_24")

	# ensure stereo
	if _arr.shape[1] == 1:
	_arr = np.repeat(_arr, 2, axis=-1)
	elif _arr.shape[1] != N_CHANNELS:
	raise gr.Error(f'{stem}: expected {N_CHANNELS} channels, got {_arr.shape[1]}')

	if _sr != SR:
	# resample to model's expected SR (44100); librosa needs float32 in [-1, 1]
	# so we convert int16 to float, resample, then convert back to int16
	temp = _arr.astype(np.float32) / (1. + np.iinfo(_arr.dtype).max)
	temp = librosa.resample(temp.T, orig_sr=_sr, target_sr=SR,
	res_type='kaiser_best').T
	temp = temp * (1 + np.iinfo(np.int16).max)
	_arr = np.clip(temp, np.iinfo(np.int16).min,
	np.iinfo(np.int16).max).astype(np.int16)

	# Pad the start with 30s of wrapped audio so the model has context at
	# t=0. This follows the resample rather than preceding it, because
	# pad_samples counts samples at SR: padding first adds that many
	# samples at the input's own rate, and resampling then leaves a pad
	# shorter than the trim after inference removes. A 48 kHz input came
	# back empty for exactly that reason.
	_arr = np.pad(_arr, [(pad_samples, 0), (0, 0)], 'wrap')

	samples.append(_arr.shape[0])
	audio[stem] = (SR, _arr)
	else:
	# missing stems are silent so the model always receives 4 tracks
	audio[stem] = (SR, np.zeros((SR, N_CHANNELS), dtype=np.int16))

	max_samples = min(max(samples), int(MAX_LENGTH * SR))
	max_samplingrate = max(samplingrate_list)
	max_subtype = max(subtype_list)

	# --- Build input tensor ---
	data = np.zeros((len(STEMS), 1, max_samples, N_CHANNELS), dtype=np.float32)
	for k, stem in enumerate(STEMS):
	stem_audio = generate_data(audio[stem])
	if stem_audio[:max_samples].shape[0] < max_samples:
	zeros_ = np.zeros_like(data[k][0])
	zeros_[:stem_audio[:max_samples].shape[0]] = stem_audio[:max_samples]
	data[k] = zeros_[:max_samples]
	else:
	data[k] = stem_audio[:max_samples]

	# round up to the next multiple of KERNEL_SIZE_ENCODER as required by the encoder
	new_samples = (1 + max_samples // KERNEL_SIZE_ENCODER) * KERNEL_SIZE_ENCODER
	data = np.pad(data, [(0, 0), (0, 0), (0, new_samples - max_samples), (0, 0)])

	# --- Effect normalization ---
	# changed from org. inference #512 because parallelization may not be safe on HF's CPU tier
	for k, stem in enumerate(STEMS):
	data[k][0] = normalize_audio_wave((data[k][0], EFFECTS, stem, features_mean))

	# --- Run model inference ---
	test_data = torch.from_numpy(data).to(DEVICE)
	with torch.no_grad():
	test_out = super_net.inference(test_data)
	audio_out = test_out[DataType.TIME_SAMPLES].cpu().numpy()
	audio_out = audio_out[..., :max_samples] # trim encoder padding
	audio_out = audio_out[..., pad_samples:] # trim the 30s context pad added at load

	# first output stem is the mix
	audio_output = audio_out[0, 0, ...]

	# --- Resample and save output ---
	if max_samplingrate != SR:
	audio_output = librosa.resample(audio_output, orig_sr=SR,
	target_sr=max_samplingrate,
	res_type='kaiser_best')

	with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f:
	out_path = f.name
	save_wav(out_path, max_samplingrate, audio_output.T, subtype=max_subtype)

	return out_path


	if __name__ == "__main__":
	with gr.Blocks(title="FxNorm Automix") as demo:
	input_components = [
	gr.Audio(type="filepath", label="Vocals").harp_required(True).set_info("The vocal stem, up to 3.5 minutes. Required."),
	gr.Audio(type="filepath", label="Bass").harp_required(False).set_info((
	"The bass stem. Optional, and left out if the song has none."
	)),
	gr.Audio(type="filepath", label="Drums").harp_required(False).set_info("The drum stem. Optional."),
	gr.Audio(type="filepath", label="Other").harp_required(False).set_info((
	"Everything else: guitars, keys, strings. Optional."
	)),
	]
	output_components = [
	gr.Audio(type="filepath", label="Mix").set_info("Automatic mix of the input stems."),
	]
	build_endpoint(
	model_card=model_card,
	input_components=input_components,
	output_components=output_components,
	process_fn=process_fn,
	)

	demo.queue().launch(show_error=True, pwa=True)