app.py

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 threading
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

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 by the background thread.
inference_pipeline = None
model_loading = True
model_error = None


def load_model():
    global inference_pipeline, model_loading, 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)
            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, since this runs in a background thread. The worker
        # moves the modules 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)
    finally:
        model_loading = False


# 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'],
)

# Start model loading in the background so the server can start immediately.
_loader = threading.Thread(target=load_model, daemon=True)
_loader.start()

# pyharp's worker imports this file to reach process_fn (as __mp_main__ rather than
# __main__), so it finishes loading here, before process_fn can reserve a GPU, instead
# of racing the check there. The server itself keeps loading in the background, which
# lets it start at once.
if __name__ != "__main__":
    _loader.join()

    # 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. The worker is the
    # process that runs process_fn, so it places them here, on its main thread, once the
    # loader is done. The server never runs process_fn, so its copy stays on the CPU.
    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 model_loading:
        raise gr.Error("Model is still loading, please wait a moment and try again.")
    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

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            accompaniment_path = f.name
        with tempfile.NamedTemporaryFile(suffix=".wav", 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)