"""
Nepali Song Voice Swap — Seed-VC on ZeroGPU — Hugging Face Space
==================================================================

Pipeline: song -> Demucs (separate vocals/instrumental) -> Seed-VC SVC
(convert vocals to your voice, F0-conditioned) -> remix over instrumental.

ZeroGPU notes:
  - All models must be moved to CUDA at module level, not inside the
    decorated function. Lazy loading into @spaces.GPU is much slower.
  - Quota is charged against the *visitor's* account by declared duration,
    so duration is computed per request from source-audio length.

Code from https://github.com/Plachtaa/seed-vc (GPL-3.0) is vendored under
./modules and ./hf_utils.py, unmodified, so upstream fixes can be dropped
in by re-copying those files.
"""

import os
import time
import tempfile
import spaces  # must be imported before torch

import gradio as gr
import numpy as np
import torch
import torchaudio
import librosa
import soundfile as sf
import yaml

from huggingface_hub import hf_hub_download

from modules.commons import build_model, load_checkpoint, recursive_munch
from hf_utils import load_custom_model_from_hf

torch.set_grad_enabled(False)
device = torch.device("cuda")

# ZeroGPU kills the function when the declared duration elapses, so a song
# long enough to overrun the declared estimate is worse than a rejected
# upload. Free-tier daily quota is ~300s of GPU total, which a 2-minute
# song at 25 steps already approaches. Raise via env var if you have quota.
MAX_SONG_SECONDS = float(os.environ.get("MAX_SONG_SECONDS", "120"))
MAX_REF_SECONDS = 25.0  # Seed-VC hard-clips reference audio beyond this

# --------------------------------------------------------------------------
# Load Seed-VC (singing, 44.1kHz, F0-conditioned) at module level.
# --------------------------------------------------------------------------

dit_checkpoint_path, dit_config_path = load_custom_model_from_hf(
    "Plachta/Seed-VC",
    "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema_v2.pth",
    "config_dit_mel_seed_uvit_whisper_base_f0_44k.yml",
)
_config = yaml.safe_load(open(dit_config_path, "r"))
_model_params = recursive_munch(_config["model_params"])
_model_params.dit_type = "DiT"
svc_model = build_model(_model_params, stage="DiT")
sr = _config["preprocess_params"]["sr"]
hop_length = _config["preprocess_params"]["spect_params"]["hop_length"]

svc_model, _, _, _ = load_checkpoint(
    svc_model, None, dit_checkpoint_path,
    load_only_params=True, ignore_modules=[], is_distributed=False,
)
for key in svc_model:
    svc_model[key].eval().to(device)
svc_model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)

from modules.campplus.DTDNN import CAMPPlus

campplus_ckpt_path = load_custom_model_from_hf(
    "funasr/campplus", "campplus_cn_common.bin", config_filename=None
)
campplus_model = CAMPPlus(feat_dim=80, embedding_size=192)
campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu"))
campplus_model.eval().to(device)

from modules.bigvgan import bigvgan

# BigVGAN.from_pretrained() is not usable with current huggingface_hub:
# the vendored PyTorchModelHubMixin subclass declares `proxies` and
# `resume_download` as *required* keyword-only args of _from_pretrained,
# but hub_mixin stopped passing either one. Downloading the two files and
# constructing the model directly is what from_pretrained does anyway, and
# it keeps the vendored file untouched.
_bigvgan_name = _model_params.vocoder.name
_bigvgan_config = hf_hub_download(repo_id=_bigvgan_name, filename="config.json")
_bigvgan_ckpt = hf_hub_download(repo_id=_bigvgan_name, filename="bigvgan_generator.pt")
bigvgan_model = bigvgan.BigVGAN(
    bigvgan.load_hparams_from_json(_bigvgan_config), use_cuda_kernel=False
)
bigvgan_model.load_state_dict(torch.load(_bigvgan_ckpt, map_location="cpu")["generator"])
bigvgan_model.remove_weight_norm()
bigvgan_model = bigvgan_model.eval().to(device)

from transformers import AutoFeatureExtractor, WhisperModel

_whisper_name = _model_params.speech_tokenizer.name
whisper_model = WhisperModel.from_pretrained(
    _whisper_name, torch_dtype=torch.float16
).to(device)
del whisper_model.decoder
whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(_whisper_name)

from modules.audio import mel_spectrogram

_mel_fn_args = {
    "n_fft": _config["preprocess_params"]["spect_params"]["n_fft"],
    "win_size": _config["preprocess_params"]["spect_params"]["win_length"],
    "hop_size": hop_length,
    "num_mels": _config["preprocess_params"]["spect_params"]["n_mels"],
    "sampling_rate": sr,
    "fmin": _config["preprocess_params"]["spect_params"].get("fmin", 0),
    # config stores fmax as the string "None"; upstream maps that to None
    "fmax": None if _config["preprocess_params"]["spect_params"].get("fmax", "None") == "None" else 8000,
    "center": False,
}
to_mel = lambda x: mel_spectrogram(x, **_mel_fn_args)

from modules.rmvpe import RMVPE

rmvpe_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None)
rmvpe = RMVPE(rmvpe_path, is_half=False, device=device)

max_context_window = sr // hop_length * 30
overlap_frame_len = 16
overlap_wave_len = overlap_frame_len * hop_length


def semantic_fn(waves_16k):
    inputs = whisper_feature_extractor(
        [waves_16k.squeeze(0).cpu().numpy()],
        return_tensors="pt", return_attention_mask=True,
    )
    input_features = whisper_model._mask_input_features(
        inputs.input_features, attention_mask=inputs.attention_mask
    ).to(device)
    outputs = whisper_model.encoder(
        input_features.to(whisper_model.encoder.dtype),
        head_mask=None, output_attentions=False,
        output_hidden_states=False, return_dict=True,
    )
    hidden = outputs.last_hidden_state.to(torch.float32)
    return hidden[:, : waves_16k.size(-1) // 320 + 1]


# --------------------------------------------------------------------------
# Load Demucs (vocal / instrumental separation) at module level.
# --------------------------------------------------------------------------

from demucs.pretrained import get_model as demucs_get_model
from demucs.apply import apply_model as demucs_apply_model

demucs_model = demucs_get_model("htdemucs")
demucs_model.eval().to(device)


def separate_vocals(song_path):
    # librosa (soundfile/audioread) rather than torchaudio.load, which is
    # deprecated in torchaudio 2.8 and whose mp3 support depends on which
    # backend the image happens to dispatch to.
    wav, _ = librosa.load(song_path, sr=demucs_model.samplerate, mono=False)
    wav = torch.from_numpy(np.atleast_2d(wav)).float()
    if wav.shape[0] == 1:
        wav = wav.repeat(2, 1)
    elif wav.shape[0] > 2:
        wav = wav[:2]

    # Mirrors demucs.api.Separator.separate_tensor: the mix stays on CPU and
    # apply_model moves each chunk to `device` itself.
    ref = wav.mean(0)
    mean, std = ref.mean(), ref.std() + 1e-8
    with torch.no_grad():
        sources = demucs_apply_model(
            demucs_model, ((wav - mean) / std)[None],
            shifts=0, split=True, overlap=0.25, device=device,
        )[0]
    sources = sources * std + mean

    vocals = sources[demucs_model.sources.index("vocals")]
    instrumental = sources.sum(0) - vocals
    return (
        vocals.mean(0).cpu().numpy(),
        instrumental.mean(0).cpu().numpy(),
        demucs_model.samplerate,
    )


# --------------------------------------------------------------------------
# Seed-VC singing voice conversion — adapted from seed-vc's app_svc.py,
# collapsed from a streaming generator to a single return value.
# --------------------------------------------------------------------------


def adjust_f0_semitones(f0, n_semitones):
    return f0 * (2 ** (n_semitones / 12))


def crossfade(chunk1, chunk2, overlap):
    fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2
    fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2
    chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out
    return chunk2


@torch.no_grad()
def convert_vocals(source_audio_np, source_sr, ref_path, diffusion_steps,
                    pitch_shift, auto_f0_adjust=True, length_adjust=1.0,
                    inference_cfg_rate=0.7):
    source_audio = librosa.resample(source_audio_np, orig_sr=source_sr, target_sr=sr) \
        if source_sr != sr else source_audio_np
    ref_audio = librosa.load(ref_path, sr=sr)[0]

    source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(device)
    ref_audio = torch.tensor(ref_audio[: sr * int(MAX_REF_SECONDS)]).unsqueeze(0).float().to(device)

    ref_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000)
    converted_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000)

    if converted_waves_16k.size(-1) <= 16000 * 30:
        S_alt = semantic_fn(converted_waves_16k)
    else:
        overlapping_time = 5
        S_alt_list = []
        buffer = None
        traversed_time = 0
        while traversed_time < converted_waves_16k.size(-1):
            if buffer is None:
                chunk = converted_waves_16k[:, traversed_time: traversed_time + 16000 * 30]
            else:
                chunk = torch.cat([
                    buffer,
                    converted_waves_16k[:, traversed_time: traversed_time + 16000 * (30 - overlapping_time)],
                ], dim=-1)
            S_chunk = semantic_fn(chunk)
            S_alt_list.append(S_chunk if traversed_time == 0 else S_chunk[:, 50 * overlapping_time:])
            buffer = chunk[:, -16000 * overlapping_time:]
            traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time
        S_alt = torch.cat(S_alt_list, dim=1)

    S_ori = semantic_fn(ref_waves_16k)

    mel = to_mel(source_audio)
    mel2 = to_mel(ref_audio)

    target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(device)
    target2_lengths = torch.LongTensor([mel2.size(2)]).to(device)

    feat2 = torchaudio.compliance.kaldi.fbank(
        ref_waves_16k, num_mel_bins=80, dither=0, sample_frequency=16000
    )
    feat2 = feat2 - feat2.mean(dim=0, keepdim=True)
    style2 = campplus_model(feat2.unsqueeze(0))

    F0_ori = rmvpe.infer_from_audio(ref_waves_16k[0], thred=0.03)
    F0_alt = rmvpe.infer_from_audio(converted_waves_16k[0], thred=0.03)
    F0_ori = torch.from_numpy(F0_ori).to(device)[None]
    F0_alt = torch.from_numpy(F0_alt).to(device)[None]

    voiced_F0_ori = F0_ori[F0_ori > 1]
    voiced_F0_alt = F0_alt[F0_alt > 1]
    log_f0_alt = torch.log(F0_alt + 1e-5)
    median_log_f0_ori = torch.median(torch.log(voiced_F0_ori + 1e-5))
    median_log_f0_alt = torch.median(torch.log(voiced_F0_alt + 1e-5))

    shifted_log_f0_alt = log_f0_alt.clone()
    if auto_f0_adjust:
        shifted_log_f0_alt[F0_alt > 1] = log_f0_alt[F0_alt > 1] - median_log_f0_alt + median_log_f0_ori
    shifted_f0_alt = torch.exp(shifted_log_f0_alt)
    if pitch_shift != 0:
        shifted_f0_alt[F0_alt > 1] = adjust_f0_semitones(shifted_f0_alt[F0_alt > 1], pitch_shift)

    cond, *_ = svc_model.length_regulator(S_alt, ylens=target_lengths, n_quantizers=3, f0=shifted_f0_alt)
    prompt_condition, *_ = svc_model.length_regulator(S_ori, ylens=target2_lengths, n_quantizers=3, f0=F0_ori)

    max_source_window = max_context_window - mel2.size(2)
    processed_frames = 0
    chunks = []
    previous_chunk = None
    while processed_frames < cond.size(1):
        chunk_cond = cond[:, processed_frames: processed_frames + max_source_window]
        is_last_chunk = processed_frames + max_source_window >= cond.size(1)
        cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1)
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            vc_target = svc_model.cfm.inference(
                cat_condition,
                torch.LongTensor([cat_condition.size(1)]).to(device),
                mel2, style2, None, diffusion_steps,
                inference_cfg_rate=inference_cfg_rate,
            )
            vc_target = vc_target[:, :, mel2.size(-1):]
        vc_wave = bigvgan_model(vc_target.float()).squeeze().cpu()
        if vc_wave.ndim == 1:
            vc_wave = vc_wave.unsqueeze(0)

        if previous_chunk is None:
            if is_last_chunk:
                chunks.append(vc_wave[0].numpy())
                break
            chunks.append(vc_wave[0, :-overlap_wave_len].numpy())
            previous_chunk = vc_wave[0, -overlap_wave_len:]
            processed_frames += vc_target.size(2) - overlap_frame_len
        elif is_last_chunk:
            chunks.append(crossfade(previous_chunk.numpy(), vc_wave[0].numpy(), overlap_wave_len))
            break
        else:
            chunks.append(crossfade(previous_chunk.numpy(), vc_wave[0, :-overlap_wave_len].numpy(), overlap_wave_len))
            previous_chunk = vc_wave[0, -overlap_wave_len:]
            processed_frames += vc_target.size(2) - overlap_frame_len

    return np.concatenate(chunks), sr


def remix(vocals, instrumental, sr_a, sr_b):
    if sr_a != sr_b:
        instrumental = librosa.resample(instrumental, orig_sr=sr_b, target_sr=sr_a)
    n = min(len(vocals), len(instrumental))
    mix = vocals[:n] + instrumental[:n]
    peak = np.abs(mix).max()
    if peak > 1.0:
        mix = mix / peak
    return mix


def _check_song(path):
    if not path:
        raise gr.Error("Upload the Nepali reference song first.")
    duration = librosa.get_duration(path=path)
    if duration > MAX_SONG_SECONDS:
        raise gr.Error(f"Song is {duration:.0f}s. Trim it to {MAX_SONG_SECONDS:.0f}s or less.")
    return path


def _check_ref(path):
    if not path:
        raise gr.Error("Upload your own voice clip first (5-15s, clean speech or singing).")
    return path


def estimate_duration(song_path, ref_path, diffusion_steps, pitch_shift):
    if not song_path:
        return 30
    seconds = librosa.get_duration(path=song_path)
    return min(int(15 + seconds * 1.3 * (int(diffusion_steps) / 25.0)), 300)


@spaces.GPU(duration=estimate_duration)
def swap_voice(song_path, ref_path, diffusion_steps, pitch_shift):
    song_path = _check_song(song_path)
    ref_path = _check_ref(ref_path)

    t0 = time.time()
    vocals_np, instrumental_np, demucs_sr = separate_vocals(song_path)
    t1 = time.time()

    converted, out_sr = convert_vocals(
        vocals_np, demucs_sr, ref_path,
        diffusion_steps=int(diffusion_steps), pitch_shift=float(pitch_shift),
    )
    t2 = time.time()

    mixed = remix(converted, instrumental_np, out_sr, demucs_sr)

    out_path = tempfile.mktemp(suffix=".wav")
    sf.write(out_path, mixed, out_sr)

    status = (
        f"Separated vocals in {t1 - t0:.1f}s, converted voice in {t2 - t1:.1f}s "
        f"({diffusion_steps} diffusion steps). Total {time.time() - t0:.1f}s."
    )
    return out_path, status


# --------------------------------------------------------------------------
# Interface
# --------------------------------------------------------------------------

CSS = """
.gradio-container { max-width: 900px !important; }
#header h1 { margin-bottom: 0.15rem; font-weight: 650; letter-spacing: -0.01em; }
#header p { margin-top: 0; opacity: 0.72; }
footer { visibility: hidden; }
"""

with gr.Blocks(title="Nepali Song Voice Swap", css=CSS, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# Nepali Song Voice Swap\n"
        "Upload a Nepali reference song and a clip of your own voice. "
        "The song's vocals are separated, re-sung in your voice, and remixed "
        "back over the original instrumental. Runs on ZeroGPU — GPU time is "
        "drawn from your own Hugging Face daily quota.",
        elem_id="header",
    )

    with gr.Row():
        with gr.Column(scale=3):
            song_in = gr.Audio(
                label="Nepali reference song (full mix, up to 4 min)",
                type="filepath", sources=["upload"],
            )
            ref_in = gr.Audio(
                label="Your voice (5-15s, clean, upload or record)",
                type="filepath", sources=["upload", "microphone"],
            )
            with gr.Row():
                steps = gr.Slider(
                    1, 100, value=25, step=1, label="Diffusion steps",
                    info="25 is a good default; 50 for best quality on singing.",
                )
                pitch = gr.Slider(
                    -24, 24, value=0, step=1, label="Pitch shift (semitones)",
                    info="Use if your voice's natural range differs a lot from the song.",
                )
            btn = gr.Button("Swap voice", variant="primary")
        with gr.Column(scale=2):
            out_audio = gr.Audio(label="Result", type="filepath")
            status = gr.Textbox(label="Timing", lines=3, interactive=False)

    btn.click(swap_voice, inputs=[song_in, ref_in, steps, pitch], outputs=[out_audio, status])

    gr.Markdown(
        "Built on [Seed-VC](https://github.com/Plachtaa/seed-vc) (GPL-3.0) and "
        "[Demucs](https://github.com/adefossez/demucs) (MIT) for source separation.\n\n"
        "Do not clone anyone's voice without their permission."
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch()