app.py

from __future__ import annotations

import os
import sys
import tempfile
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Tuple

import copy
import wave
import base64

try:
    import spaces
except ImportError:
    class _SpacesShim:
        @staticmethod
        def GPU(*_args, **_kwargs):
            def decorator(fn):
                return fn
            return decorator

    spaces = _SpacesShim()

import gradio as gr
import torch
from huggingface_hub import hf_hub_download, snapshot_download


APP_ROOT = Path(__file__).resolve().parent
HEARTLIB_SRC = APP_ROOT / "heartlib" / "src"
HEARTLIB_PACKAGE_INIT = HEARTLIB_SRC / "heartlib" / "__init__.py"
if HEARTLIB_PACKAGE_INIT.is_file():
    sys.path.insert(0, str(HEARTLIB_SRC))

from heartlib import HeartMuLaGenPipeline


@dataclass(frozen=True)
class ModelConfig:
    version: str
    generator_repo: str
    mula_repo: str
    codec_repo: str
    mula_dirname: str
    codec_dirname: str


MODEL_CONFIG = ModelConfig(
    version=os.getenv("HEARTMULA_VERSION", "3B"),
    generator_repo=os.getenv("HEARTMULA_GENERATOR_REPO", "HeartMuLa/HeartMuLaGen"),
    mula_repo=os.getenv(
        "HEARTMULA_MULA_REPO", "HeartMuLa/HeartMuLa-oss-3B-happy-new-year"
    ),
    codec_repo=os.getenv(
        "HEARTMULA_CODEC_REPO", "HeartMuLa/HeartCodec-oss-20260123"
    ),
    mula_dirname=os.getenv("HEARTMULA_MULA_DIRNAME", "HeartMuLa-oss-3B"),
    codec_dirname=os.getenv("HEARTMULA_CODEC_DIRNAME", "HeartCodec-oss"),
)

GPU_DURATION_SECONDS = 100
COMPILE_DURATION_SECONDS = 100
MAX_DURATION_SECONDS = 120
DEFAULT_DURATION_SECONDS = min(
    int(os.getenv("HEARTMULA_DEFAULT_DURATION_SECONDS", "60")),
    MAX_DURATION_SECONDS,
)
ENABLE_FLASH_ATTN = os.getenv("HEARTMULA_ENABLE_FLASH_ATTN", "1") != "0"
ENABLE_AOTI = os.getenv("HEARTMULA_ENABLE_AOTI", "1") != "0"
AOTI_MAX_BATCH = int(os.getenv("HEARTMULA_AOTI_MAX_BATCH", "2"))
AOTI_MAX_SEQ_LEN = int(os.getenv("HEARTMULA_AOTI_MAX_SEQ_LEN", "4096"))
KEEP_MULA_LOADED = os.getenv("HEARTMULA_KEEP_MULA_LOADED", "1") != "0"
KEEP_CODEC_LOADED = os.getenv("HEARTMULA_KEEP_CODEC_LOADED", "0") != "0"

MODEL_LOCK = threading.Lock()
PIPELINE_LOCK = threading.Lock()
PIPELINE_CACHE: Dict[Tuple[str, str], HeartMuLaGenPipeline] = {}
_RUNTIME_PREPARED = False


def _default_cache_root() -> Path:
    env_home = os.getenv("HF_HOME")
    if env_home:
        return Path(env_home)

    data_home = Path("/data/.huggingface")
    if data_home.parent.exists():
        return data_home

    return Path("/tmp/huggingface")


def _model_root() -> Path:
    return Path(
        os.getenv(
            "HEARTMULA_MODEL_DIR",
            str(_default_cache_root() / "heartmula_models"),
        )
    )


def _read_text(path: Path, fallback: str) -> str:
    if path.is_file():
        return path.read_text(encoding="utf-8").strip()
    return fallback


def _cached_model_exists(model_dir: Path) -> bool:
    required_paths = [
        model_dir / "tokenizer.json",
        model_dir / "gen_config.json",
        model_dir / MODEL_CONFIG.mula_dirname,
        model_dir / MODEL_CONFIG.codec_dirname,
    ]
    return all(path.exists() for path in required_paths)


def ensure_model_artifacts(progress: gr.Progress | None = None) -> Path:
    model_dir = _model_root()
    model_dir.mkdir(parents=True, exist_ok=True)

    if _cached_model_exists(model_dir):
        if progress is not None:
            progress(0.05, desc="Using cached model artifacts")
        return model_dir

    with MODEL_LOCK:
        if _cached_model_exists(model_dir):
            if progress is not None:
                progress(0.05, desc="Using cached model artifacts")
            return model_dir

        if progress is not None:
            progress(0.05, desc="Downloading tokenizer and generation config")
        for filename in ("tokenizer.json", "gen_config.json"):
            hf_hub_download(
                repo_id=MODEL_CONFIG.generator_repo,
                filename=filename,
                local_dir=str(model_dir),
            )

        if progress is not None:
            progress(0.25, desc="Downloading HeartMuLa checkpoint")
        snapshot_download(
            repo_id=MODEL_CONFIG.mula_repo,
            local_dir=str(model_dir / MODEL_CONFIG.mula_dirname),
        )

        if progress is not None:
            progress(0.6, desc="Downloading HeartCodec checkpoint")
        snapshot_download(
            repo_id=MODEL_CONFIG.codec_repo,
            local_dir=str(model_dir / MODEL_CONFIG.codec_dirname),
        )

    if progress is not None:
        progress(0.95, desc="Model artifacts ready")
    return model_dir


def get_pipeline(model_dir: Path) -> HeartMuLaGenPipeline:
    runtime = "cuda" if torch.cuda.is_available() else "cpu"
    cache_key = (runtime, str(model_dir))

    with PIPELINE_LOCK:
        if cache_key in PIPELINE_CACHE:
            return PIPELINE_CACHE[cache_key]

        if runtime == "cuda":
            device = {
                "mula": torch.device("cuda"),
                "codec": torch.device("cuda"),
            }
            dtype = {
                "mula": torch.bfloat16,
                "codec": torch.float32,
            }
            lazy_load = {
                "mula": not KEEP_MULA_LOADED,
                "codec": not KEEP_CODEC_LOADED,
            }
        else:
            device = torch.device("cpu")
            dtype = torch.float32
            lazy_load = False

        pipeline = HeartMuLaGenPipeline.from_pretrained(
            str(model_dir),
            device=device,
            dtype=dtype,
            version=MODEL_CONFIG.version,
            lazy_load=lazy_load,
        )
        pipeline.configure_runtime_acceleration(
            enable_flash_attn=runtime == "cuda" and ENABLE_FLASH_ATTN,
            enable_aoti=runtime == "cuda" and ENABLE_AOTI,
            max_batch_size=AOTI_MAX_BATCH,
            max_compile_seq_len=AOTI_MAX_SEQ_LEN,
        )
        PIPELINE_CACHE[cache_key] = pipeline
        return pipeline


@spaces.GPU(duration=COMPILE_DURATION_SECONDS)
def _compile_runtime(model_dir: str):
    pipeline = get_pipeline(Path(model_dir))
    pipeline.prepare_runtime()


def apply_voice_conversion(
    audio_path: str,
    engine: str,
    model_path: str | None,
    index_path: str | None,
    f0_up_key: int,
    f0_method: str,
    index_rate: float,
    protect: float,
) -> str:
    if engine == "Ninguno":
        return audio_path

    try:
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
            output_path = fp.name

        if engine == "RVC":
            if not model_path:
                return audio_path
            from rvc_python.infer import infer_file
            infer_file(
                input_path=audio_path,
                output_path=output_path,
                model_path=model_path,
                index_path=index_path if index_path else "",
                f0_up_key=f0_up_key,
                f0_method=f0_method,
                index_rate=index_rate,
                protect=protect,
            )
            return output_path

        elif engine == "So-VITS-SVC":
            if not model_path:
                return audio_path
            from svclib.infer import infer_file as svc_infer
            svc_infer(
                input_path=audio_path,
                output_path=output_path,
                model_path=model_path,
                config_path=index_path if index_path else "",
                tran=f0_up_key,
                cluster_path="",
            )
            return output_path

        elif engine == "Coqui TTS (XTTS)":
            if not model_path:
                return audio_path
            from TTS.api import TTS
            tts = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2", progress_bar=False).to("cuda" if torch.cuda.is_available() else "cpu")
            tts.voice_conversion_to_file(
                source_wav=audio_path,
                target_wav=model_path,
                output_path=output_path
            )
            return output_path

        elif engine == "ACE Studio / ACE Step 1.5":
            from acestep.infer import infer_file as ace_infer
            ace_infer(
                input_path=audio_path,
                output_path=output_path,
                model_path=model_path if model_path else "",
                pitch_shift=f0_up_key,
            )
            return output_path

    except Exception:
        return audio_path

    return audio_path


@spaces.GPU(duration=GPU_DURATION_SECONDS)
def _run_generation(
    model_dir: str,
    lyrics: str,
    tags: str,
    max_duration_seconds: int,
    temperature: float,
    topk: int,
    cfg_scale: float,
    engine: str,
    model_file: str | None,
    index_file: str | None,
    f0_up_key: int,
    f0_method: str,
    index_rate: float,
    protect: float,
    progress=gr.Progress(track_tqdm=True),
):
    pipeline = get_pipeline(Path(model_dir))
    max_audio_length_ms = max_duration_seconds * 1000

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
        output_path = fp.name

    progress(0.05, desc="Generating audio")
    with torch.no_grad():
        pipeline(
            {
                "lyrics": lyrics,
                "tags": tags,
            },
            max_audio_length_ms=max_audio_length_ms,
            save_path=output_path,
            topk=topk,
            temperature=temperature,
            cfg_scale=cfg_scale,
        )

    if engine != "Ninguno":
        progress(0.85, desc=f"Applying {engine} Conversion")
        output_path = apply_voice_conversion(
            audio_path=output_path,
            engine=engine,
            model_path=model_file,
            index_path=index_file,
            f0_up_key=f0_up_key,
            f0_method=f0_method,
            index_rate=index_rate,
            protect=protect,
        )

    with open(output_path, "rb") as f:
        audio_bytes = f.read()
        
    audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
    return output_path, audio_base64


def generate_music(
    lyrics: str,
    tags: str,
    max_duration_seconds: int,
    temperature: float,
    topk: int,
    cfg_scale: float,
    engine: str,
    model_upload,
    index_upload,
    f0_up_key: int,
    f0_method: str,
    index_rate: float,
    protect: float,
    progress=gr.Progress(track_tqdm=True),
):
    if not lyrics.strip():
        raise gr.Error("Please enter lyrics before generating.")
    if not tags.strip():
        raise gr.Error("Please enter at least one style tag.")

    model_dir = ensure_model_artifacts(progress)

    global _RUNTIME_PREPARED
    if not _RUNTIME_PREPARED:
        progress(0.02, desc="Compiling runtime (first request, please wait)")
        _compile_runtime(str(model_dir))
        _RUNTIME_PREPARED = True

    model_path = model_upload.name if hasattr(model_upload, "name") else model_upload
    index_path = index_upload.name if hasattr(index_upload, "name") else index_upload

    return _run_generation(
        str(model_dir),
        lyrics,
        tags,
        max_duration_seconds,
        temperature,
        topk,
        cfg_scale,
        engine,
        model_path,
        index_path,
        f0_up_key,
        f0_method,
        index_rate,
        protect,
    )


@spaces.GPU(duration=300)
def train_rvc_model(
    audio_file,
    model_name: str,
    sample_rate: int,
    f0_method: str,
    progress=gr.Progress(track_tqdm=True),
):
    if not audio_file:
        raise gr.Error("Please upload a WAV audio file for training.")
    if not model_name.strip():
        raise gr.Error("Please enter a model name.")

    audio_path = audio_file.name if hasattr(audio_file, "name") else audio_file

    try:
        progress(0.1, desc="Preparing dataset and slicing audio")
        from rvc_python.train import train_pipeline

        output_dir = tempfile.mkdtemp()
        
        progress(0.3, desc="Extracting features (F0 & Embeddings)")
        progress(0.6, desc="Training model for 100 epochs")
        
        pth_path, index_path = train_pipeline(
            experiment_name=model_name.strip(),
            dataset_path=audio_path,
            sr=sample_rate,
            f0_method=f0_method,
            total_epoch=100,
            save_every_epoch=100,
            gpus="0" if torch.cuda.is_available() else "",
            output_dir=output_dir,
        )

        progress(1.0, desc="Training complete!")
        return pth_path, index_path

    except Exception as e:
        raise gr.Error(f"Training failed: {str(e)}")


DEFAULT_LYRICS = _read_text(
    APP_ROOT / "heartlib" / "assets" / "lyrics.txt",
    """[Verse]
The city wakes before the sun
We keep moving one by one

[Chorus]
Hold the light and sing it through
Every road comes back to you""",
)

DEFAULT_TAGS = _read_text(
    APP_ROOT / "heartlib" / "assets" / "tags.txt",
    "female,indie pop,piano,emotional,night,silky,memories",
)


with gr.Blocks(title="HeartMuLa + Multi-Engine Voice Cloning & Training") as demo:
    gr.Markdown(
        """
        # HeartMuLa + Multi-Engine Voice Cloning & RVC Trainer
        """
    )

    with gr.Tabs():
        with gr.TabItem("Generación de Música y Clonación"):
            with gr.Row():
                with gr.Column(scale=1):
                    lyrics_input = gr.Textbox(
                        label="Lyrics",
                        lines=14,
                        value=DEFAULT_LYRICS,
                        placeholder="Use structured sections such as [Verse], [Chorus], [Bridge].",
                    )
                    tags_input = gr.Textbox(
                        label="Tags",
                        value=DEFAULT_TAGS,
                        placeholder="female,indie pop,piano,emotional,night,silky,memories",
                        info="Comma-separated tags without spaces for best compatibility.",
                    )
                    with gr.Accordion("Voice Cloning Settings", open=True):
                        engine_select = gr.Radio(
                            choices=["Ninguno", "RVC", "So-VITS-SVC", "Coqui TTS (XTTS)", "ACE Studio / ACE Step 1.5"],
                            value="Ninguno",
                            label="Cloning Engine"
                        )
                        model_upload = gr.File(label="Model File / Checkpoint / Target WAV")
                        index_upload = gr.File(label="Index / Config File [Optional]")
                        
                        f0_up_key = gr.Slider(
                            minimum=-24, maximum=24, value=0, step=1,
                            label="Pitch Shift / Semitones"
                        )
                        f0_method = gr.Dropdown(
                            choices=["rmvpe", "pm", "harvest", "crepe"],
                            value="rmvpe",
                            label="Pitch Extraction Method (RVC)"
                        )
                        index_rate = gr.Slider(
                            minimum=0.0, maximum=1.0, value=0.75, step=0.05,
                            label="Index Rate (RVC)"
                        )
                        protect = gr.Slider(
                            minimum=0.0, maximum=0.5, value=0.33, step=0.01,
                            label="Protect Vocals/Breaths (RVC)"
                        )

                    with gr.Accordion("Generation Settings", open=False):
                        max_duration_input = gr.Slider(
                            minimum=30,
                            maximum=MAX_DURATION_SECONDS,
                            value=DEFAULT_DURATION_SECONDS,
                            step=10,
                            label="Max Duration (seconds)",
                        )
                        temperature_input = gr.Slider(
                            minimum=0.1,
                            maximum=2.0,
                            value=1.0,
                            step=0.1,
                            label="Temperature",
                        )
                        topk_input = gr.Slider(
                            minimum=1,
                            maximum=100,
                            value=50,
                            step=1,
                            label="Top-K",
                        )
                        cfg_scale_input = gr.Slider(
                            minimum=1.0,
                            maximum=3.0,
                            value=1.5,
                            step=0.1,
                            label="CFG Scale",
                        )

                    generate_button = gr.Button("Generate Music", variant="primary")

                with gr.Column(scale=1):
                    audio_output = gr.Audio(label="Generated Audio", type="filepath")
                    output_json = gr.Textbox(interactive=False, show_label=False, info='Tokens used to generate the audio.')

        with gr.TabItem("Entrenar Modelo RVC (100 Epochs)"):
            with gr.Row():
                with gr.Column(scale=1):
                    train_audio_input = gr.File(label="Audio Dataset (.wav)")
                    train_name_input = gr.Textbox(label="Model Name", value="my_voice_model")
                    train_sr_select = gr.Dropdown(choices=[32000, 40000, 48000], value=40000, label="Sample Rate")
                    train_f0_select = gr.Dropdown(choices=["rmvpe", "pm", "harvest", "crepe"], value="rmvpe", label="Pitch Extraction Method")
                    
                    train_button = gr.Button("Start Training (100 Epochs)", variant="primary")

                with gr.Column(scale=1):
                    output_pth = gr.File(label="Generated .pth Weight File")
                    output_index = gr.File(label="Generated .index File")

    generate_button.click(
        fn=generate_music,
        inputs=[
            lyrics_input,
            tags_input,
            max_duration_input,
            temperature_input,
            topk_input,
            cfg_scale_input,
            engine_select,
            model_upload,
            index_upload,
            f0_up_key,
            f0_method,
            index_rate,
            protect,
        ],
        outputs=[audio_output, output_json]
    )

    train_button.click(
        fn=train_rvc_model,
        inputs=[
            train_audio_input,
            train_name_input,
            train_sr_select,
            train_f0_select,
        ],
        outputs=[output_pth, output_index]
    )

demo.queue(default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch()