"""
MelodyFlow: edits music to follow a text description while keeping its timing and structure.

audiocraft is vendored from the authors' Space (facebook/MelodyFlow), whose fork moves it to the
torch release ZeroGPU requires, and runs here rather than being called over the network.
"""

import spaces  # imported before torch, as ZeroGPU requires

import os
import tempfile

import gradio as gr
import torch

# audiocraft's checkpoints pickle their omegaconf configuration, which torch 2.6 and newer refuse to
# load by default (weights_only=True). The authors' Space restores the old default the same way, and
# only Meta's own checkpoints are loaded here. Patched before audiocraft is imported.
_torch_load = torch.load


def _load_audiocraft_checkpoint(*args, **kwargs):
    kwargs.setdefault("weights_only", False)
    return _torch_load(*args, **kwargs)


torch.load = _load_audiocraft_checkpoint

from pyharp import *
from audiocraft.data.audio import audio_read, audio_write
from audiocraft.data.audio_utils import convert_audio
from audiocraft.models import MelodyFlow

MODEL_NAME = "facebook/melodyflow-t24-30secs"
TAKES = 3
# The rate and channel count the model's compression model was trained on
MODEL_SR, MODEL_CHANNELS = 48000, 2

# Fetched when the Space builds (preload_from_hub in README.md)
model = MelodyFlow.get_pretrained(MODEL_NAME)


model_card = ModelCard(
    name="MelodyFlow",
    description=(
        "Edits music to follow a text description while keeping the timing and structure of the "
        "recording it was given. Flow matching carries the input part of the way along a path "
        "and regenerates from there, so a low target flow step stays close to the original and a "
        "high one commits to the prompt. With no input audio it generates from the text alone, up "
        "to thirty seconds. The code is MIT, but the MelodyFlow weights are CC-BY-NC-4.0, so "
        "results cannot be used commercially. Repository: "
        "https://github.com/facebookresearch/audiocraft Paper: https://arxiv.org/abs/2407.03648"
    ),
    author="Gael Le Lan, Bowen Shi, Zhaoheng Ni, et al.",
    tags=[Subcategory.EDITING],
)


@spaces.GPU(duration=90)
@torch.inference_mode()
def process_fn(text, steps, target_flowstep, regularize, regularization_strength, duration, melody):
    steps = int(steps)
    if melody is not None:
        # Editing inverts the input first, which the authors' demo pays for with half the steps
        steps = steps // 2

    model.set_generation_params(solver="midpoint", steps=steps, duration=duration)
    model.set_editing_params(solver="midpoint", steps=steps, target_flowstep=target_flowstep,
                             regularize=regularize, lambda_kl=regularization_strength)

    texts = [text] * TAKES
    if melody is not None:
        wav, sr = audio_read(melody)
        if wav.dim() == 2:
            wav = wav[None]
        wav = wav[..., : int(sr * model.duration)]
        wav = convert_audio(wav, sr, MODEL_SR, MODEL_CHANNELS)
        tokens = model.encode_audio(wav.to(model.device))
        outputs = model.edit(prompt_tokens=tokens.repeat(TAKES, 1, 1), descriptions=texts,
                             src_descriptions=[""] * TAKES, return_tokens=False)
    else:
        outputs = model.generate(texts, return_tokens=False)

    out_dir = tempfile.mkdtemp()
    paths = []
    for i, output in enumerate(outputs.detach().cpu().float(), start=1):
        path = os.path.join(out_dir, f"melodyflow_take_{i}.wav")
        # The loudness normalization the authors' demo applies
        audio_write(path, output, model.sample_rate, strategy="loudness",
                    loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)
        paths.append(path)
    return tuple(paths)


# The processing worker imports this file, so the app must not be built there
if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Textbox(label="Input Text",
            info=(
                "Describe what the music should become. The timing and structure of "
                "the input are kept, so this changes the character rather than the "
                "arrangement."
            )),
            gr.Slider(minimum=1, maximum=250, step=1, value=128.0, label="Inference steps",
            info="More steps give a cleaner result and take longer."),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.0, label="Target Flow step",
            info=(
                "How far the edit travels from the input. Low values stay close to "
                "the original, high values commit to the prompt."
            )),
            gr.Checkbox(value=False, label="Regularize",
            info=(
                "Keeps the result closer to the input's structure. Turn it off for "
                "a freer edit."
            )),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.2, label="Regularization Strength",
            info=(
                "How strongly that structure is held. Only used when Regularize is "
                "on."
            )),
            gr.Slider(minimum=1, maximum=30, step=1, value=30.0, label="Duration",
            info="Length of the result, in seconds."),
            gr.Audio(type="filepath", label="File or Microphone").harp_required(False).set_info(
                "The music to edit. Leave it empty to generate from the text alone."
            ),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Variation 1").set_info((
                "First of three takes from the same prompt and settings."
            )),
            gr.Audio(type="filepath", label="Variation 2").set_info((
                "Second take."
            )),
            gr.Audio(type="filepath", label="Variation 3").set_info("Third take."),
        ]
        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)