Spaces:

teamup-tech
/

AnyAccomp

Running on Zero

Drop the comment on ssr_mode

15dc4ff verified 1 day ago

8.71 kB

import spaces  # imported before torch, as ZeroGPU requires
import sys
import numpy as np
sys.stdout.reconfigure(line_buffering=True) # ensures print() appears in HF Spaces logs in real time

import traceback
import contextlib
import tempfile
import torch
import os
import librosa
import soundfile as sf
import gradio as gr
from huggingface_hub import snapshot_download
from pyharp import ModelCard, build_endpoint, SampleRate, Channels
from pyharp.tags import Subcategory

from anyaccomp.inference_utils import Sing2SongInferencePipeline

# Create a ModelCard describing this endpoint for HARP.
model_card = ModelCard(
    name="AnyAccomp",
    description=(
        "Upload a 3-30 second vocal or instrument track and the model will "
        "generate an accompaniment. The code is MIT and the weights are "
        "CC-BY-4.0, which requires attribution. Generation runs at 24 kHz in "
        "mono. Both outputs are resampled back to the input's rate and copied "
        "across its channels, so they match the format they came in with but "
        "carry no stereo image. Repository: "
        "https://github.com/AmphionTeam/AnyAccomp Paper: "
        "https://arxiv.org/abs/2509.14052"
    ),
    author="Junan Zhang, Yunjia Zhang, Xueyao Zhang, Zhizheng Wu",
    tags=[Subcategory.ACCOMPANIMENT, SampleRate(24000), Channels(1), 'music', 'domain:music'],
)

# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the result
# lossily again would expose the artifacts of the first encode and add a second set.
LOSSLESS = {".wav", ".flac"}

repo_id = "amphion/anyaccomp"
base_dir = os.path.dirname(os.path.abspath(__file__))
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Model loading state, populated in the worker
inference_pipeline = None
model_error = None


def load_model():
    global inference_pipeline, model_error
    try:
        checkpoint_marker = os.path.join(base_dir, "pretrained", "flow_matching")
        if not os.path.exists(checkpoint_marker):
            print(f"Downloading model files from {repo_id}...", flush=True)
            model_dir = snapshot_download(repo_id=repo_id, local_dir=base_dir,
                                          revision="9aa9e62427337bf1df4caa3c4f3e6ad934522e71")
            print(f"Model files downloaded to: {model_dir}", flush=True)
        else:
            model_dir = base_dir
            print("Model files already present, skipping download.", flush=True)

        cfg_path = os.path.join(base_dir, "config/flow_matching.json")
        vocoder_cfg_path = os.path.join(base_dir, "config/vocoder.json")
        checkpoint_path = os.path.join(model_dir, "pretrained/flow_matching")
        vocoder_checkpoint_path = os.path.join(model_dir, "pretrained/vocoder")

        print("Initializing AnyAccomp InferencePipeline...", flush=True)
        # Loaded on the CPU, and moved to the GPU once loading is done (see below)
        pipeline = Sing2SongInferencePipeline(
            checkpoint_path, cfg_path,
            vocoder_checkpoint_path, vocoder_cfg_path,
            device="cpu",
        )
        pipeline.sample_rate = 24000
        inference_pipeline = pipeline
        print("Model loaded successfully.", flush=True)
    except Exception as e:
        model_error = str(e)
        print(f"Error loading model: {e}", flush=True)


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ (rather than
# __main__) and starts with the server, so the model is downloaded and loaded there alone, before
# process_fn can reserve a GPU. A failure is kept rather than raised, so that each request can
# report it. ZeroGPU wants models placed on CUDA at module level, where its emulation records each
# placement for every GPU call to replay, rather than inside the call, so the worker places them
# here, on its main thread, once loading is done.
if __name__ != "__main__":
    load_model()
    if inference_pipeline is not None:
        for module in (inference_pipeline.model, inference_pipeline.coco_model,
                       inference_pipeline.vocoder):
            module.to(DEVICE)
        inference_pipeline.device = DEVICE


def process_fn(vocal_filepath, n_timesteps):  # add , cfg_scale to re-enable
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if inference_pipeline is None:
        raise gr.Error(f"Model failed to load: {model_error}")

    if vocal_filepath is None:
        raise gr.Error("Please upload a vocal audio file.")

    try:
        duration = librosa.get_duration(path=vocal_filepath)
    except Exception as e:
        raise gr.Error(f"Cannot read the audio file: {e}")
    if not (3 <= duration <= 30):
        raise gr.Error(f"Audio must be 3 to 30 seconds long, not {duration:.1f}.")

    return accompany(vocal_filepath, n_timesteps)


@spaces.GPU(duration=60)
@torch.inference_mode()
def accompany(vocal_filepath, n_timesteps):
    try:
        # The session's format, read before the downmix and resample the model
        # needs, and restored on both outputs
        info = sf.info(vocal_filepath)
        original_sr = info.samplerate
        original_channels = info.channels

        vocal_audio, _ = librosa.load(vocal_filepath, sr=24000, mono=True)
        vocal_tensor = torch.tensor(vocal_audio).unsqueeze(0).to(DEVICE)

        vocal_mel = inference_pipeline.encode_vocal(vocal_tensor)

        autocast_ctx = torch.amp.autocast("cuda", dtype=torch.bfloat16) if torch.cuda.is_available() else contextlib.nullcontext()
        with autocast_ctx:
            mel = inference_pipeline.model.reverse_diffusion(
                vocal_mel=vocal_mel,
                n_timesteps=int(n_timesteps),
                cfg=3.0,  # replace with cfg_scale to re-enable
            )

        mel = mel.float()
        wav = inference_pipeline._generate_audio(mel)
        wav = wav.squeeze().detach().cpu().numpy()

        wav = librosa.util.fix_length(data=wav, size=len(vocal_audio))
        mixture_wav = wav + vocal_audio

        input_ext = os.path.splitext(vocal_filepath)[1].lower()
        output_ext = input_ext if input_ext in LOSSLESS else ".wav"
        with tempfile.NamedTemporaryFile(suffix=output_ext, delete=False) as f:
            accompaniment_path = f.name
        with tempfile.NamedTemporaryFile(suffix=output_ext, delete=False) as f:
            mixture_path = f.name

        def _restore(audio):
            # Generation is mono at 24 kHz, so the channels are copies rather
            # than a stereo image, which the model card states
            if original_sr != 24000:
                audio = librosa.resample(
                    audio, orig_sr=24000, target_sr=original_sr
                )
            if original_channels > 1:
                audio = np.tile(audio[:, None], (1, original_channels))
            return audio

        sf.write(accompaniment_path, _restore(wav), original_sr)
        sf.write(mixture_path, _restore(mixture_wav), original_sr)

        return accompaniment_path, mixture_path

    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"An error occurred during processing: {e}")


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(
                type="filepath",
                label="Input Vocal or Instrument Audio",
            ).harp_required(True).set_info((
                "A dry solo vocal or lead line, 3 to 30 seconds. Accompaniment is "
                "written around what it hears, so a clean take gives a better fit."
            )),
            gr.Slider(
                minimum=10, maximum=100, value=50, step=1,
                label="Inference Steps (n_timesteps)",
                info="Number of diffusion steps. More steps = slower but higher quality",
            ),
            # gr.Slider(minimum=1.0, maximum=10.0, value=3.0, step=0.1, label="CFG Scale", info="Guidance strength — higher values follow the input melody more closely."),
        ]

        output_components = [
            gr.Audio(type="filepath", label="Generated Accompaniment").set_info(
                "The generated instrumental accompaniment."
            ),
            gr.Audio(type="filepath", label="Mixture (Vocal + Accompaniment)").set_info(
                "Vocal mixed with the generated accompaniment."
            ),
        ]

        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )


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

	import spaces # imported before torch, as ZeroGPU requires
	import sys
	import numpy as np
	sys.stdout.reconfigure(line_buffering=True) # ensures print() appears in HF Spaces logs in real time

	import traceback
	import contextlib
	import tempfile
	import torch
	import os
	import librosa
	import soundfile as sf
	import gradio as gr
	from huggingface_hub import snapshot_download
	from pyharp import ModelCard, build_endpoint, SampleRate, Channels
	from pyharp.tags import Subcategory

	from anyaccomp.inference_utils import Sing2SongInferencePipeline

	# Create a ModelCard describing this endpoint for HARP.
	model_card = ModelCard(
	name="AnyAccomp",
	description=(
	"Upload a 3-30 second vocal or instrument track and the model will "
	"generate an accompaniment. The code is MIT and the weights are "
	"CC-BY-4.0, which requires attribution. Generation runs at 24 kHz in "
	"mono. Both outputs are resampled back to the input's rate and copied "
	"across its channels, so they match the format they came in with but "
	"carry no stereo image. Repository: "
	"https://github.com/AmphionTeam/AnyAccomp Paper: "
	"https://arxiv.org/abs/2509.14052"
	),
	author="Junan Zhang, Yunjia Zhang, Xueyao Zhang, Zhizheng Wu",
	tags=[Subcategory.ACCOMPANIMENT, SampleRate(24000), Channels(1), 'music', 'domain:music'],
	)

	# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the result
	# lossily again would expose the artifacts of the first encode and add a second set.
	LOSSLESS = {".wav", ".flac"}

	repo_id = "amphion/anyaccomp"
	base_dir = os.path.dirname(os.path.abspath(__file__))
	DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

	# Model loading state, populated in the worker
	inference_pipeline = None
	model_error = None


	def load_model():
	global inference_pipeline, model_error
	try:
	checkpoint_marker = os.path.join(base_dir, "pretrained", "flow_matching")
	if not os.path.exists(checkpoint_marker):
	print(f"Downloading model files from {repo_id}...", flush=True)
	model_dir = snapshot_download(repo_id=repo_id, local_dir=base_dir,
	revision="9aa9e62427337bf1df4caa3c4f3e6ad934522e71")
	print(f"Model files downloaded to: {model_dir}", flush=True)
	else:
	model_dir = base_dir
	print("Model files already present, skipping download.", flush=True)

	cfg_path = os.path.join(base_dir, "config/flow_matching.json")
	vocoder_cfg_path = os.path.join(base_dir, "config/vocoder.json")
	checkpoint_path = os.path.join(model_dir, "pretrained/flow_matching")
	vocoder_checkpoint_path = os.path.join(model_dir, "pretrained/vocoder")

	print("Initializing AnyAccomp InferencePipeline...", flush=True)
	# Loaded on the CPU, and moved to the GPU once loading is done (see below)
	pipeline = Sing2SongInferencePipeline(
	checkpoint_path, cfg_path,
	vocoder_checkpoint_path, vocoder_cfg_path,
	device="cpu",
	)
	pipeline.sample_rate = 24000
	inference_pipeline = pipeline
	print("Model loaded successfully.", flush=True)
	except Exception as e:
	model_error = str(e)
	print(f"Error loading model: {e}", flush=True)


	# pyharp runs process_fn in a worker, which imports this file as __mp_main__ (rather than
	# __main__) and starts with the server, so the model is downloaded and loaded there alone, before
	# process_fn can reserve a GPU. A failure is kept rather than raised, so that each request can
	# report it. ZeroGPU wants models placed on CUDA at module level, where its emulation records each
	# placement for every GPU call to replay, rather than inside the call, so the worker places them
	# here, on its main thread, once loading is done.
	if __name__ != "__main__":
	load_model()
	if inference_pipeline is not None:
	for module in (inference_pipeline.model, inference_pipeline.coco_model,
	inference_pipeline.vocoder):
	module.to(DEVICE)
	inference_pipeline.device = DEVICE


	def process_fn(vocal_filepath, n_timesteps): # add , cfg_scale to re-enable
	# Checked before the GPU is reserved, so a request the model cannot take spends none of the
	# caller's ZeroGPU allowance
	if inference_pipeline is None:
	raise gr.Error(f"Model failed to load: {model_error}")

	if vocal_filepath is None:
	raise gr.Error("Please upload a vocal audio file.")

	try:
	duration = librosa.get_duration(path=vocal_filepath)
	except Exception as e:
	raise gr.Error(f"Cannot read the audio file: {e}")
	if not (3 <= duration <= 30):
	raise gr.Error(f"Audio must be 3 to 30 seconds long, not {duration:.1f}.")

	return accompany(vocal_filepath, n_timesteps)


	@spaces.GPU(duration=60)
	@torch.inference_mode()
	def accompany(vocal_filepath, n_timesteps):
	try:
	# The session's format, read before the downmix and resample the model
	# needs, and restored on both outputs
	info = sf.info(vocal_filepath)
	original_sr = info.samplerate
	original_channels = info.channels

	vocal_audio, _ = librosa.load(vocal_filepath, sr=24000, mono=True)
	vocal_tensor = torch.tensor(vocal_audio).unsqueeze(0).to(DEVICE)

	vocal_mel = inference_pipeline.encode_vocal(vocal_tensor)

	autocast_ctx = torch.amp.autocast("cuda", dtype=torch.bfloat16) if torch.cuda.is_available() else contextlib.nullcontext()
	with autocast_ctx:
	mel = inference_pipeline.model.reverse_diffusion(
	vocal_mel=vocal_mel,
	n_timesteps=int(n_timesteps),
	cfg=3.0, # replace with cfg_scale to re-enable
	)

	mel = mel.float()
	wav = inference_pipeline._generate_audio(mel)
	wav = wav.squeeze().detach().cpu().numpy()

	wav = librosa.util.fix_length(data=wav, size=len(vocal_audio))
	mixture_wav = wav + vocal_audio

	input_ext = os.path.splitext(vocal_filepath)[1].lower()
	output_ext = input_ext if input_ext in LOSSLESS else ".wav"
	with tempfile.NamedTemporaryFile(suffix=output_ext, delete=False) as f:
	accompaniment_path = f.name
	with tempfile.NamedTemporaryFile(suffix=output_ext, delete=False) as f:
	mixture_path = f.name

	def _restore(audio):
	# Generation is mono at 24 kHz, so the channels are copies rather
	# than a stereo image, which the model card states
	if original_sr != 24000:
	audio = librosa.resample(
	audio, orig_sr=24000, target_sr=original_sr
	)
	if original_channels > 1:
	audio = np.tile(audio[:, None], (1, original_channels))
	return audio

	sf.write(accompaniment_path, _restore(wav), original_sr)
	sf.write(mixture_path, _restore(mixture_wav), original_sr)

	return accompaniment_path, mixture_path

	except Exception as e:
	traceback.print_exc()
	raise gr.Error(f"An error occurred during processing: {e}")


	if __name__ == "__main__":
	with gr.Blocks() as demo:
	input_components = [
	gr.Audio(
	type="filepath",
	label="Input Vocal or Instrument Audio",
	).harp_required(True).set_info((
	"A dry solo vocal or lead line, 3 to 30 seconds. Accompaniment is "
	"written around what it hears, so a clean take gives a better fit."
	)),
	gr.Slider(
	minimum=10, maximum=100, value=50, step=1,
	label="Inference Steps (n_timesteps)",
	info="Number of diffusion steps. More steps = slower but higher quality",
	),
	# gr.Slider(minimum=1.0, maximum=10.0, value=3.0, step=0.1, label="CFG Scale", info="Guidance strength — higher values follow the input melody more closely."),
	]

	output_components = [
	gr.Audio(type="filepath", label="Generated Accompaniment").set_info(
	"The generated instrumental accompaniment."
	),
	gr.Audio(type="filepath", label="Mixture (Vocal + Accompaniment)").set_info(
	"Vocal mixed with the generated accompaniment."
	),
	]

	build_endpoint(
	model_card=model_card,
	input_components=input_components,
	output_components=output_components,
	process_fn=process_fn,
	)


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