core.py

"""
Core audio processing module for Hyper-RVC.

Provides functions for:
- Model and music downloading
- Audio separation (vocals, instrumentals, karaoke)
- RVC voice conversion
- Audio effects (reverb, denoise, deecho)
- Audio merging and post-processing
"""

import sys
import os
import subprocess
import torch
from functools import lru_cache
import shutil
from typing import Optional, Dict, Any, List, Tuple
from pedalboard import Pedalboard, Reverb
from pedalboard.io import AudioFile
from pydub import AudioSegment
from audio_separator.separator import Separator
import logging
import yaml
import gc

now_dir = os.getcwd()
sys.path.append(now_dir)

from programs.applio_code.rvc.infer.infer import VoiceConverter
from programs.applio_code.rvc.lib.tools.model_download import model_download_pipeline
from programs.music_separation_code.inference import proc_file
from programs.variable import (
    models_vocals,
    karaoke_models,
    denoise_models,
    deecho_models,
    dereverb_models,
    check_fp16_support
)
from programs.tools.logger import get_logger

logger = get_logger(__name__)


@lru_cache(maxsize=None)
def import_voice_converter() -> VoiceConverter:
    """
    Import and cache the VoiceConverter instance.

    Returns:
        VoiceConverter instance for RVC inference
    """
    from programs.applio_code.rvc.infer.infer import VoiceConverter
    return VoiceConverter()


@lru_cache(maxsize=1)
def get_config():
    """
    Get and cache the RVC configuration.

    Returns:
        RVC Config instance
    """
    from programs.applio_code.rvc.configs.config import Config
    return Config()


def download_file(url: str, path: str, filename: str) -> Optional[str]:
    """
    Download a file from a URL to a specified path.

    Args:
        url: URL to download from
        path: Directory path to save the file
        filename: Name of the file to save

    Returns:
        Path to downloaded file if successful, None otherwise
    """
    os.makedirs(path, exist_ok=True)
    file_path = os.path.join(path, filename)

    if os.path.exists(file_path):
        logger.info(f"File '{filename}' already exists at '{path}'.")
        return file_path

    try:
        torch.hub.download_url_to_file(url, file_path)
        logger.info(f"File '{filename}' downloaded successfully")
        return file_path
    except Exception as e:
        logger.error(f"Error downloading file '{filename}' from '{url}': {e}")
        return None


def get_model_info_by_name(model_name: str) -> Optional[Dict[str, Any]]:
    """
    Get model information by model name.

    Args:
        model_name: Name of the model to search for

    Returns:
        Dictionary containing model information, or None if not found
    """
    all_models = (
        models_vocals
        + karaoke_models
        + dereverb_models
        + deecho_models
        + denoise_models
    )
    for model in all_models:
        if model["name"] == model_name:
            return model
    return None


def get_last_modified_file(pasta: str) -> Optional[str]:
    """
    Get the most recently modified file in a directory.

    Args:
        pasta: Directory path to search

    Returns:
        Name of the most recently modified file, or None if directory is empty

    Raises:
        NotADirectoryError: If the path is not a valid directory
    """
    if not os.path.isdir(pasta):
        raise NotADirectoryError(f"{pasta} is not a valid directory.")
    arquivos = [f for f in os.listdir(pasta) if os.path.isfile(os.path.join(pasta, f))]
    if not arquivos:
        return None
    return max(arquivos, key=lambda x: os.path.getmtime(os.path.join(pasta, x)))


def search_with_word(folder: str, word: str) -> Optional[str]:
    """
    Search for the most recent file containing a specific word in its name.

    Args:
        folder: Directory to search in
        word: Word to search for in filenames

    Returns:
        Name of the most recent file containing the word, or None if not found

    Raises:
        NotADirectoryError: If the path is not a valid directory
    """
    if not os.path.isdir(folder):
        raise NotADirectoryError(f"{folder} is not a valid directory.")
    file_with_word = [file for file in os.listdir(folder) if word in file]
    if not file_with_word:
        return None
    most_recent_file = max(
        file_with_word, key=lambda file: os.path.getmtime(os.path.join(folder, file))
    )
    return most_recent_file


def search_with_two_words(folder: str, word1: str, word2: str) -> Optional[str]:
    """
    Search for the most recent file containing two specific words in its name.

    Args:
        folder: Directory to search in
        word1: First word to search for
        word2: Second word to search for

    Returns:
        Name of the most recent file containing both words, or None if not found

    Raises:
        NotADirectoryError: If the path is not a valid directory
    """
    if not os.path.isdir(folder):
        raise NotADirectoryError(f"{folder} is not a valid directory.")
    file_with_words = [
        file for file in os.listdir(folder) if word1 in file and word2 in file
    ]
    if not file_with_words:
        return None
    most_recent_file = max(
        file_with_words, key=lambda file: os.path.getmtime(os.path.join(folder, file))
    )
    return most_recent_file


def get_last_modified_folder(path: str) -> Optional[str]:
    """
    Get the most recently modified subdirectory.

    Args:
        path: Parent directory path

    Returns:
        Path to the most recently modified subdirectory, or None if none exist
    """
    directories = [
        os.path.join(path, d)
        for d in os.listdir(path)
        if os.path.isdir(os.path.join(path, d))
    ]
    if not directories:
        return None
    last_modified_folder = max(directories, key=os.path.getmtime)
    return last_modified_folder


def download_model(link: str) -> str:
    """
    Download an RVC model from a link.

    Args:
        link: URL or link to download the model from

    Returns:
        Success or error message
    """
    try:
        model_download_pipeline(link)
        logger.info("Model downloaded successfully")
        return "Model downloaded with success"
    except Exception as e:
        logger.error(f"Error downloading model: {e}")
        return f"Error downloading model: {e}"


def download_music(link: str) -> str:
    """
    Download music from a URL (YouTube, etc.).

    Args:
        link: URL to download music from

    Returns:
        Success or error message
    """
    try:
        os.makedirs(os.path.join(now_dir, "audio_files", "original_files"), exist_ok=True)
        command = [
            "yt-dlp",
            "-x",
            "--output",
            os.path.join(now_dir, "audio_files", "original_files", "%(title)s.%(ext)s"),
            link,
        ]
        subprocess.run(command, check=True)
        logger.info("Music downloaded successfully")
        return "Music downloaded with success"
    except subprocess.CalledProcessError as e:
        logger.error(f"Error downloading music: {e}")
        return f"Error downloading music: {e}"
    except Exception as e:
        logger.error(f"Unexpected error downloading music: {e}")
        return f"Error downloading music: {e}"


def whisper_process(
    model_size: str,
    input_audio: str,
    configs: Dict[str, Any],
    device: str,
    out_queue,
    word_timestamps: bool = True
) -> None:
    """
    Process audio with Whisper for transcription/speaker diarization.

    Args:
        model_size: Size of the Whisper model to use
        input_audio: Path to input audio file
        configs: Configuration dictionary
        device: Device to run inference on
        out_queue: Queue to put results in
        word_timestamps: Whether to extract word-level timestamps
    """
    from programs.speaker_diarization.whisper import load_model

    try:
        segments = load_model(
            model_size,
            device=device
        ).transcribe(
            input_audio,
            fp16=check_fp16_support(device),
            word_timestamps=word_timestamps
        )

        out_queue.put(segments["segments"])
    except Exception as e:
        out_queue.put(e)
    finally:
        del segments
        gc.collect()


def add_audio_effects(
    audio_path: str,
    reverb_size: float,
    reverb_wet: float,
    reverb_dry: float,
    reverb_damping: float,
    reverb_width: float,
    output_path: str,
) -> str:
    """
    Add reverb effects to an audio file using Pedalboard.

    Args:
        audio_path: Path to input audio file
        reverb_size: Room size (0-1)
        reverb_wet: Wet level (0-1)
        reverb_dry: Dry level (0-1)
        reverb_damping: Damping (0-1)
        reverb_width: Stereo width (0-1)
        output_path: Path to save the output file

    Returns:
        Path to the output file
    """
    try:
        board = Pedalboard([])
        board.append(
            Reverb(
                room_size=reverb_size,
                dry_level=reverb_dry,
                wet_level=reverb_wet,
                damping=reverb_damping,
                width=reverb_width,
            )
        )
        with AudioFile(audio_path) as f:
            with AudioFile(output_path, "w", f.samplerate, f.num_channels) as o:
                while f.tell() < f.frames:
                    chunk = f.read(int(f.samplerate))
                    effected = board(chunk, f.samplerate, reset=False)
                    o.write(effected)
        logger.info(f"Audio effects applied successfully: {output_path}")
        return output_path
    except Exception as e:
        logger.error(f"Error applying audio effects: {e}")
        raise


def merge_audios(
    vocals_path: str,
    inst_path: str,
    backing_path: str,
    output_path: str,
    main_gain: float,
    inst_gain: float,
    backing_Vol: float,
    output_format: str,
) -> str:
    """
    Merge multiple audio files (vocals, instrumentals, backing vocals).

    Args:
        vocals_path: Path to vocals audio file
        inst_path: Path to instrumental audio file
        backing_path: Path to backing vocals audio file
        output_path: Path to save the merged output
        main_gain: Volume gain for main vocals (dB)
        inst_gain: Volume gain for instrumentals (dB)
        backing_Vol: Volume gain for backing vocals (dB)
        output_format: Output format (e.g., 'mp3', 'flac', 'wav')

    Returns:
        Path to the merged output file
    """
    try:
        main_vocal_audio = AudioSegment.from_file(vocals_path, format="flac") + main_gain
        instrumental_audio = AudioSegment.from_file(inst_path, format="flac") + inst_gain
        backing_vocal_audio = (
            AudioSegment.from_file(backing_path, format="flac") + backing_Vol
        )
        combined_audio = main_vocal_audio.overlay(
            instrumental_audio.overlay(backing_vocal_audio)
        )
        combined_audio.export(output_path, format=output_format)
        logger.info(f"Audio files merged successfully: {output_path}")
        return output_path
    except Exception as e:
        logger.error(f"Error merging audio files: {e}")
        raise


def update_model_config_for_fp16(model_info: Dict[str, Any], use_fp16: bool) -> None:
    """
    Update model config file based on FP16 support.

    Args:
        model_info: Dictionary containing model information
        use_fp16: Whether FP16 is supported
    """
    if not use_fp16 and os.path.exists(model_info["config"]):
        try:
            with open(model_info["config"], "r") as file:
                config = yaml.safe_load(file)

            if "training" in config and "use_amp" in config["training"]:
                config["training"]["use_amp"] = False

            with open(model_info["config"], "w") as file:
                yaml.safe_dump(config, file)
            logger.info(f"Disabled FP16/AMP in config for {model_info['name']}")
        except Exception as e:
            logger.error(f"Error updating config for {model_info['name']}: {e}")


def full_inference_program(
    model_path: str,
    index_path: str,
    input_audio_path: str,
    output_path: str,
    export_format_rvc: str,
    split_audio: bool,
    autotune: bool,
    vocal_model: str,
    karaoke_model: str,
    dereverb_model: str,
    deecho: bool,
    deecho_model: str,
    denoise: bool,
    denoise_model: str,
    reverb: bool,
    vocals_volume: float,
    instrumentals_volume: float,
    backing_vocals_volume: float,
    export_format_final: str,
    devices: str,
    pitch: int,
    filter_radius: int,
    index_rate: float,
    rms_mix_rate: float,
    protect: float,
    pitch_extract: str,
    hop_lenght: int,
    reverb_room_size: float,
    reverb_damping: float,
    reverb_wet_gain: float,
    reverb_dry_gain: float,
    reverb_width: float,
    embedder_model: str,
    delete_audios: bool,
    use_tta: bool,
    batch_size: int,
    infer_backing_vocals: bool,
    infer_backing_vocals_model: str,
    infer_backing_vocals_index: str,
    change_inst_pitch: int,
    pitch_back: int,
    filter_radius_back: int,
    index_rate_back: float,
    rms_mix_rate_back: float,
    protect_back: float,
    pitch_extract_back: str,
    hop_length_back: int,
    export_format_rvc_back: str,
    split_audio_back: bool,
    autotune_back: bool,
    embedder_model_back: str,
) -> Tuple[str, str]:
    """
    Run the full RVC inference pipeline on an audio file.

    This function performs:
    1. Vocal separation from the input audio
    2. Karaoke/backing vocal separation
    3. Dereverb processing (optional)
    4. Deecho processing (optional)
    5. Denoise processing (optional)
    6. RVC voice conversion
    7. Backing vocals inference (optional)
    8. Reverb effects (optional)
    9. Pitch adjustment for instrumentals (optional)
    10. Final audio merging

    Args:
        model_path: Path to the RVC model file
        index_path: Path to the RVC index file
        input_audio_path: Path to the input audio file
        output_path: Directory path for output
        export_format_rvc: Export format for RVC output
        split_audio: Whether to split audio for processing
        autotune: Whether to apply autotune
        vocal_model: Name of the vocal separation model
        karaoke_model: Name of the karaoke separation model
        dereverb_model: Name of the dereverb model
        deecho: Whether to apply deecho processing
        deecho_model: Name of the deecho model
        denoise: Whether to apply denoise processing
        denoise_model: Name of the denoise model
        reverb: Whether to apply reverb effect
        vocals_volume: Volume adjustment for vocals (dB)
        instrumentals_volume: Volume adjustment for instrumentals (dB)
        backing_vocals_volume: Volume adjustment for backing vocals (dB)
        export_format_final: Final export format
        devices: GPU devices to use (e.g., "0" or "0 1")
        pitch: Pitch adjustment in semitones
        filter_radius: Filter radius for pitch extraction
        index_rate: Index rate for voice conversion
        rms_mix_rate: RMS mix rate for volume envelope
        protect: Protect value for conversion
        pitch_extract: Pitch extraction method
        hop_lenght: Hop length for pitch extraction
        reverb_room_size: Reverb room size (0-1)
        reverb_damping: Reverb damping (0-1)
        reverb_wet_gain: Reverb wet gain (0-1)
        reverb_dry_gain: Reverb dry gain (0-1)
        reverb_width: Reverb width (0-1)
        embedder_model: Embedder model to use
        delete_audios: Whether to delete intermediate audio files
        use_tta: Whether to use test-time augmentation
        batch_size: Batch size for separation models
        infer_backing_vocals: Whether to infer backing vocals separately
        infer_backing_vocals_model: Model for backing vocals inference
        infer_backing_vocals_index: Index for backing vocals inference
        change_inst_pitch: Change instrumental pitch (semitones)
        pitch_back: Pitch for backing vocals
        filter_radius_back: Filter radius for backing vocals
        index_rate_back: Index rate for backing vocals
        rms_mix_rate_back: RMS mix rate for backing vocals
        protect_back: Protect value for backing vocals
        pitch_extract_back: Pitch extraction for backing vocals
        hop_length_back: Hop length for backing vocals
        export_format_rvc_back: Export format for backing vocals
        split_audio_back: Whether to split audio for backing vocals
        autotune_back: Whether to apply autotune to backing vocals
        embedder_model_back: Embedder model for backing vocals

    Returns:
        Tuple of (success message, output file path)

    Raises:
        Exception: If any step in the pipeline fails
    """
    # Initialize configuration
    configs = {}

    # Determine device and FP16 support
    if torch.cuda.is_available() and devices != "cpu":
        n_gpu = torch.cuda.device_count()
        devices = devices.replace("-", " ")
        logger.info(f"Number of GPUs available: {n_gpu}")
        first_device = devices.split()[0] if devices.split() else "cuda:0"

        # Check config for FP16 setting, fallback to hardware check
        use_fp16 = configs.get("fp16", check_fp16_support(first_device))
        logger.info(f"FP16 inference: {'Enabled' if use_fp16 else 'Disabled'}")
    else:
        devices = "cpu"
        logger.info("Using CPU")
        use_fp16 = False

    music_folder = os.path.splitext(os.path.basename(input_audio_path))[0]

    # Vocals Separation
    model_info = get_model_info_by_name(vocal_model)
    model_ckpt_path = os.path.join(model_info["path"], "model.ckpt")
    if not os.path.exists(model_ckpt_path):
        download_file(
            model_info["model_url"],
            model_info["path"],
            "model.ckpt",
        )
    config_json_path = os.path.join(model_info["path"], "config.yaml")
    if not os.path.exists(config_json_path):
        download_file(
            model_info["config_url"],
            model_info["path"],
            "config.yaml",
        )

    # Update model config based on FP16 support
    update_model_config_for_fp16(model_info, use_fp16)

    store_dir = os.path.join(now_dir, "audio_files", music_folder, "vocals")
    inst_dir = os.path.join(now_dir, "audio_files", music_folder, "instrumentals")
    os.makedirs(store_dir, exist_ok=True)
    os.makedirs(inst_dir, exist_ok=True)
    input_audio_basename = os.path.splitext(os.path.basename(input_audio_path))[0]
    search_result = search_with_word(store_dir, "vocals")
    if search_result:
        logger.info("Vocals already separated")
    else:
        logger.info("Separating vocals")
        command = [
            "python",
            os.path.join(now_dir, "programs", "music_separation_code", "inference.py"),
            "--model_type",
            model_info["type"],
            "--config_path",
            model_info["config"],
            "--start_check_point",
            model_info["model"],
            "--input_file",
            input_audio_path,
            "--store_dir",
            store_dir,
            "--flac_file",
            "--pcm_type",
            "PCM_16",
            "--extract_instrumental",
        ]

        if devices == "cpu":
            command.append("--force_cpu")
        else:
            device_ids = [str(int(device)) for device in devices.split()]
            command.extend(["--device_ids"] + device_ids)

        # Add FP16 flag if supported
        if use_fp16:
            command.append("--fp16")

        subprocess.run(command)
        os.rename(
            os.path.join(
                store_dir,
                search_with_two_words(
                    store_dir,
                    os.path.basename(input_audio_path).split(".")[0],
                    "instrumental",
                ),
            ),
            os.path.join(
                inst_dir,
                f"{os.path.basename(input_audio_path).split('.')[0]}_instrumentals.flac",
            ),
        )
    inst_file = os.path.join(
        inst_dir,
        search_with_two_words(
            inst_dir, os.path.basename(input_audio_path).split(".")[0], "instrumentals"
        ),
    )

    # karaoke separation
    model_info = get_model_info_by_name(karaoke_model)
    store_dir = os.path.join(now_dir, "audio_files", music_folder, "karaoke")
    os.makedirs(store_dir, exist_ok=True)
    vocals_path = os.path.join(now_dir, "audio_files", music_folder, "vocals")
    input_file = search_with_word(vocals_path, "vocals")
    karaoke_exists = search_with_word(store_dir, "karaoke") is not None

    if karaoke_exists:
        logger.info("Backing vocals already separated")
    else:
        if input_file:
            input_file = os.path.join(vocals_path, input_file)
        logger.info("Separating backing vocals")
        if model_info["name"] == "Mel-Roformer Karaoke by aufr33 and viperx":
            model_ckpt_path = os.path.join(model_info["path"], "model.ckpt")
            if not os.path.exists(model_ckpt_path):
                download_file(
                    model_info["model_url"],
                    model_info["path"],
                    "model.ckpt",
                )
            config_json_path = os.path.join(model_info["path"], "config.yaml")
            if not os.path.exists(config_json_path):
                download_file(
                    model_info["config_url"],
                    model_info["path"],
                    "config.yaml",
                )
            
            # Update model config based on FP16 support
            update_model_config_for_fp16(model_info, use_fp16)

            command = [
                "python",
                os.path.join(
                    now_dir, "programs", "music_separation_code", "inference.py"
                ),
                "--model_type",
                model_info["type"],
                "--config_path",
                model_info["config"],
                "--start_check_point",
                model_info["model"],
                "--input_file",
                input_file,
                "--store_dir",
                store_dir,
                "--flac_file",
                "--pcm_type",
                "PCM_16",
                "--extract_instrumental",
            ]

            if devices == "cpu":
                command.append("--force_cpu")
            else:
                device_ids = [str(int(device)) for device in devices.split()]
                command.extend(["--device_ids"] + device_ids)

            # Add FP16 flag if supported
            if use_fp16:
                command.append("--fp16")

            subprocess.run(command)
        else:
            vr_params = {
                "batch_size": batch_size,
                "enable_tta": use_tta,
            }
            if use_fp16:
                vr_params["fp16"] = True
                
            separator = Separator(
                model_file_dir=os.path.join(now_dir, "models", "karaoke"),
                log_level=logging.WARNING,
                normalization_threshold=1.0,
                output_format="flac",
                output_dir=store_dir,
                vr_params=vr_params,
            )
            separator.load_model(model_filename=model_info["full_name"])
            separator.separate(input_file)
            karaoke_path = os.path.join(now_dir, "audio_files", music_folder, "karaoke")
            vocals_result = search_with_two_words(
                karaoke_path,
                os.path.basename(input_audio_path).split(".")[0],
                "Vocals",
            )
            instrumental_result = search_with_two_words(
                karaoke_path,
                os.path.basename(input_audio_path).split(".")[0],
                "Instrumental",
            )
            if "UVR-BVE-4B_SN-44100-1" in os.path.basename(vocals_result):
                os.rename(
                    os.path.join(karaoke_path, vocals_result),
                    os.path.join(
                        karaoke_path,
                        f"{os.path.basename(input_audio_path).split('.')[0]}_karaoke.flac",
                    ),
                )
            if "UVR-BVE-4B_SN-44100-1" in os.path.basename(instrumental_result):
                os.rename(
                    os.path.join(karaoke_path, instrumental_result),
                    os.path.join(
                        karaoke_path,
                        f"{os.path.basename(input_audio_path).split('.')[0]}_instrumental.flac",
                    ),
                )

    # dereverb
    model_info = get_model_info_by_name(dereverb_model)
    store_dir = os.path.join(now_dir, "audio_files", music_folder, "dereverb")
    os.makedirs(store_dir, exist_ok=True)
    karaoke_path = os.path.join(now_dir, "audio_files", music_folder, "karaoke")
    input_file = search_with_word(karaoke_path, "karaoke")
    noreverb_exists = search_with_word(store_dir, "noreverb") is not None
    if noreverb_exists:
        logger.info("Reverb already removed")
    else:
        if input_file:
            input_file = os.path.join(karaoke_path, input_file)
        logger.info("Removing reverb")
        if (
            model_info["name"] == "BS-Roformer Dereverb by anvuew"
            or model_info["name"] == "MDX23C DeReverb by aufr33 and jarredou"
        ):
            model_ckpt_path = os.path.join(model_info["path"], "model.ckpt")
            if not os.path.exists(model_ckpt_path):
                download_file(
                    model_info["model_url"],
                    model_info["path"],
                    "model.ckpt",
                )
            config_json_path = os.path.join(model_info["path"], "config.yaml")
            if not os.path.exists(config_json_path):
                download_file(
                    model_info["config_url"],
                    model_info["path"],
                    "config.yaml",
                )
            
            # Update model config based on FP16 support
            update_model_config_for_fp16(model_info, use_fp16)
            
            command = [
                "python",
                os.path.join(
                    now_dir, "programs", "music_separation_code", "inference.py"
                ),
                "--model_type",
                model_info["type"],
                "--config_path",
                model_info["config"],
                "--start_check_point",
                model_info["model"],
                "--input_file",
                input_file,
                "--store_dir",
                store_dir,
                "--flac_file",
                "--pcm_type",
                "PCM_16",
            ]

            if devices == "cpu":
                command.append("--force_cpu")
            else:
                device_ids = [str(int(device)) for device in devices.split()]
                command.extend(["--device_ids"] + device_ids)

            # Add FP16 flag if supported
            if use_fp16:
                command.append("--fp16")

            subprocess.run(command)
        else:
            if model_info["arch"] == "vr":
                vr_params = {
                    "batch_size": batch_size,
                    "enable_tta": use_tta,
                }
                if use_fp16:
                    vr_params["fp16"] = True
                separator = Separator(
                    model_file_dir=os.path.join(now_dir, "models", "dereverb"),
                    log_level=logging.WARNING,
                    normalization_threshold=1.0,
                    output_format="flac",
                    output_dir=store_dir,
                    output_single_stem="No Reverb",
                    vr_params=vr_params,
                )
            else:
                separator = Separator(
                    model_file_dir=os.path.join(now_dir, "models", "dereverb"),
                    log_level=logging.WARNING,
                    normalization_threshold=1.0,
                    output_format="flac",
                    output_dir=store_dir,
                    output_single_stem="No Reverb",
                )
            separator.load_model(model_filename=model_info["full_name"])
            separator.separate(input_file)
            dereverb_path = os.path.join(
                now_dir, "audio_files", music_folder, "dereverb"
            )
            search_result = search_with_two_words(
                dereverb_path,
                os.path.basename(input_audio_path).split(".")[0],
                "No Reverb",
            )
            if "UVR-DeEcho-DeReverb" in os.path.basename(
                search_result
            ) or "MDX Reverb HQ by FoxJoy" in os.path.basename(search_result):
                os.rename(
                    os.path.join(dereverb_path, search_result),
                    os.path.join(
                        dereverb_path,
                        f"{os.path.basename(input_audio_path).split('.')[0]}_noreverb.flac",
                    ),
                )

    # deecho
    store_dir = os.path.join(now_dir, "audio_files", music_folder, "deecho")
    os.makedirs(store_dir, exist_ok=True)
    if deecho:
        no_echo_exists = search_with_word(store_dir, "noecho") is not None
        if no_echo_exists:
            logger.info("Echo already removed")
        else:
            logger.info("Removing echo")
            model_info = get_model_info_by_name(deecho_model)

            dereverb_path = os.path.join(
                now_dir, "audio_files", music_folder, "dereverb"
            )
            noreverb_file = search_with_word(dereverb_path, "noreverb")

            input_file = os.path.join(dereverb_path, noreverb_file)

            vr_params = {
                "batch_size": batch_size,
                "enable_tta": use_tta,
            }
            if use_fp16:
                vr_params["fp16"] = True
                
            separator = Separator(
                model_file_dir=os.path.join(now_dir, "models", "deecho"),
                log_level=logging.WARNING,
                normalization_threshold=1.0,
                output_format="flac",
                output_dir=store_dir,
                output_single_stem="No Echo",
                vr_params=vr_params,
            )
            separator.load_model(model_filename=model_info["full_name"])
            separator.separate(input_file)
            deecho_path = os.path.join(now_dir, "audio_files", music_folder, "deecho")
            search_result = search_with_two_words(
                deecho_path,
                os.path.basename(input_audio_path).split(".")[0],
                "No Echo",
            )
            if "UVR-De-Echo-Normal" in os.path.basename(
                search_result
            ) or "UVR-Deecho-Agggressive" in os.path.basename(search_result):
                os.rename(
                    os.path.join(deecho_path, search_result),
                    os.path.join(
                        deecho_path,
                        f"{os.path.basename(input_audio_path).split('.')[0]}_noecho.flac",
                    ),
                )

    # denoise
    store_dir = os.path.join(now_dir, "audio_files", music_folder, "denoise")
    os.makedirs(store_dir, exist_ok=True)
    if denoise:
        no_noise_exists = search_with_word(store_dir, "dry") is not None
        if no_noise_exists:
            logger.info("Noise already removed")
        else:
            model_info = get_model_info_by_name(denoise_model)
            logger.info("Removing noise")
            input_file = (
                os.path.join(
                    now_dir,
                    "audio_files",
                    music_folder,
                    "deecho",
                    search_with_word(
                        os.path.join(now_dir, "audio_files", music_folder, "deecho"),
                        "noecho",
                    ),
                )
                if deecho
                else os.path.join(
                    now_dir,
                    "audio_files",
                    music_folder,
                    "dereverb",
                    search_with_word(
                        os.path.join(now_dir, "audio_files", music_folder, "dereverb"),
                        "noreverb",
                    ),
                )
            )

            if (
                model_info["name"] == "Mel-Roformer Denoise Normal by aufr33"
                or model_info["name"] == "Mel-Roformer Denoise Aggressive by aufr33"
            ):
                model_ckpt_path = os.path.join(model_info["path"], "model.ckpt")
                if not os.path.exists(model_ckpt_path):
                    download_file(
                        model_info["model_url"],
                        model_info["path"],
                        "model.ckpt",
                    )
                config_json_path = os.path.join(model_info["path"], "config.yaml")
                if not os.path.exists(config_json_path):
                    download_file(
                        model_info["config_url"], model_info["path"], "config.yaml"
                    )
                
                # Update model config based on FP16 support
                update_model_config_for_fp16(model_info, use_fp16)
                
                command = [
                    "python",
                    os.path.join(
                        now_dir, "programs", "music_separation_code", "inference.py"
                    ),
                    "--model_type",
                    model_info["type"],
                    "--config_path",
                    model_info["config"],
                    "--start_check_point",
                    model_info["model"],
                    "--input_file",
                    input_file,
                    "--store_dir",
                    store_dir,
                    "--flac_file",
                    "--pcm_type",
                    "PCM_16",
                ]

                if devices == "cpu":
                    command.append("--force_cpu")
                else:
                    device_ids = [str(int(device)) for device in devices.split()]
                    command.extend(["--device_ids"] + device_ids)

                # Add FP16 flag if supported
                if use_fp16:
                    command.append("--fp16")

                subprocess.run(command)
            else:
                vr_params = {
                    "batch_size": batch_size,
                    "enable_tta": use_tta,
                }
                if use_fp16:
                    vr_params["fp16"] = True
                    
                separator = Separator(
                    model_file_dir=os.path.join(now_dir, "models", "denoise"),
                    log_level=logging.WARNING,
                    normalization_threshold=1.0,
                    output_format="flac",
                    output_dir=store_dir,
                    output_single_stem="No Noise",
                    vr_params=vr_params,
                )
                separator.load_model(model_filename=model_info["full_name"])
                separator.separate(input_file)
                denoise_path = os.path.join(now_dir, "audio_files", music_folder, "denoise")
                search_result = search_with_two_words(
                    denoise_path,
                    os.path.basename(input_audio_path).split(".")[0],
                    "No Noise",
                )
                if "UVR Denoise" in os.path.basename(search_result):
                    os.rename(
                        os.path.join(denoise_path, search_result),
                        os.path.join(
                            denoise_path,
                            f"{os.path.basename(input_audio_path).split('.')[0]}_dry.flac",
                        ),
                    )

    # RVC
    denoise_path = os.path.join(now_dir, "audio_files", music_folder, "denoise")
    deecho_path = os.path.join(now_dir, "audio_files", music_folder, "deecho")
    dereverb_path = os.path.join(now_dir, "audio_files", music_folder, "dereverb")

    denoise_audio = search_with_two_words(
        denoise_path, os.path.basename(input_audio_path).split(".")[0], "dry"
    )
    deecho_audio = search_with_two_words(
        deecho_path, os.path.basename(input_audio_path).split(".")[0], "noecho"
    )
    dereverb = search_with_two_words(
        dereverb_path, os.path.basename(input_audio_path).split(".")[0], "noreverb"
    )

    if denoise_audio:
        final_path = os.path.join(
            now_dir, "audio_files", music_folder, "denoise", denoise_audio
        )
    elif deecho_audio:
        final_path = os.path.join(
            now_dir, "audio_files", music_folder, "deecho", deecho_audio
        )
    elif dereverb:
        final_path = os.path.join(
            now_dir, "audio_files", music_folder, "dereverb", dereverb
        )
    else:
        final_path = None

    store_dir = os.path.join(now_dir, "audio_files", music_folder, "rvc")
    os.makedirs(store_dir, exist_ok=True)
    logger.info("Making RVC inference")
    output_rvc = os.path.join(
        now_dir,
        "audio_files",
        music_folder,
        "rvc",
        f"{os.path.basename(input_audio_path).split('.')[0]}_rvc.wav",
    )
    inference_vc = import_voice_converter()
    inference_vc.convert_audio(
        audio_input_path=final_path,
        audio_output_path=output_rvc,
        model_path=model_path,
        index_path=index_path,
        embedder_model=embedder_model,
        pitch=pitch,
        f0_file=None,
        f0_method=pitch_extract,
        filter_radius=filter_radius,
        index_rate=index_rate,
        volume_envelope=rms_mix_rate,
        protect=protect,
        split_audio=split_audio,
        f0_autotune=autotune,
        hop_length=hop_lenght,
        export_format=export_format_rvc,
        embedder_model_custom=None,
    )
    backing_vocals = os.path.join(
        karaoke_path, search_with_word(karaoke_path, "instrumental")
    )

    if infer_backing_vocals:
        logger.info("Inferring backing vocals")
        karaoke_path = os.path.join(now_dir, "audio_files", music_folder, "karaoke")
        instrumental_file = search_with_word(karaoke_path, "instrumental")
        backing_vocals = os.path.join(karaoke_path, instrumental_file)
        output_backing_vocals = os.path.join(
            karaoke_path, f"{input_audio_basename}_instrumental_output.wav"
        )
        inference_vc.convert_audio(
            audio_input_path=backing_vocals,
            audio_output_path=output_backing_vocals,
            model_path=infer_backing_vocals_model,
            index_path=infer_backing_vocals_index,
            embedder_model=embedder_model_back,
            pitch=pitch_back,
            f0_file=None,
            f0_method=pitch_extract_back,
            filter_radius=filter_radius_back,
            index_rate=index_rate_back,
            volume_envelope=rms_mix_rate_back,
            protect=protect_back,
            split_audio=split_audio_back,
            f0_autotune=autotune_back,
            hop_length=hop_length_back,
            export_format=export_format_rvc_back,
            embedder_model_custom=None,
        )
        backing_vocals = output_backing_vocals

    # post process
    if reverb:
        add_audio_effects(
            os.path.join(
                now_dir,
                "audio_files",
                music_folder,
                "rvc",
                get_last_modified_file(
                    os.path.join(now_dir, "audio_files", music_folder, "rvc")
                ),
            ),
            reverb_room_size,
            reverb_wet_gain,
            reverb_dry_gain,
            reverb_damping,
            reverb_width,
            os.path.join(
                now_dir,
                "audio_files",
                music_folder,
                "rvc",
                os.path.basename(input_audio_path),
            ),
        )
    if change_inst_pitch != 0:
        logger.info("Changing instrumental pitch")
        inst_path = os.path.join(
            now_dir,
            "audio_files",
            music_folder,
            "instrumentals",
            search_with_word(
                os.path.join(now_dir, "audio_files", music_folder, "instrumentals"),
                "instrumentals",
            ),
        )
        audio = AudioSegment.from_file(inst_path)

        factor = 2 ** (change_inst_pitch / 12)

        new_frame_rate = int(audio.frame_rate * factor)
        audio = audio._spawn(audio.raw_data, overrides={"frame_rate": new_frame_rate})

        audio = audio.set_frame_rate(audio.frame_rate)
        output_dir_pitch = os.path.join(
            now_dir, "audio_files", music_folder, "instrumentals"
        )
        output_path_pitch = os.path.join(
            output_dir_pitch, "inst_with_changed_pitch.flac"
        )
        audio.export(output_path_pitch, format="flac")

    # merge audios
    store_dir = os.path.join(now_dir, "audio_files", music_folder, "final")
    os.makedirs(store_dir, exist_ok=True)

    vocals_path = os.path.join(now_dir, "audio_files", music_folder, "rvc")
    vocals_file = get_last_modified_file(
        os.path.join(now_dir, "audio_files", music_folder, "rvc")
    )
    vocals_file = os.path.join(vocals_path, vocals_file)

    karaoke_path = os.path.join(now_dir, "audio_files", music_folder, "karaoke")
    karaoke_file = search_with_word(karaoke_path, "Instrumental") or search_with_word(
        karaoke_path, "instrumental"
    )
    karaoke_file = os.path.join(karaoke_path, karaoke_file)
    final_output_path = os.path.join(
        now_dir,
        "audio_files",
        music_folder,
        "final",
        f"{os.path.basename(input_audio_path).split('.')[0]}_final.{export_format_final.lower()}",
    )
    logger.info("Merging audios")
    result = merge_audios(
        vocals_file,
        inst_file,
        backing_vocals,
        final_output_path,
        vocals_volume,
        instrumentals_volume,
        backing_vocals_volume,
        export_format_final,
    )
    logger.info("Audios merged!")
    if delete_audios:
        main_directory = os.path.join(now_dir, "audio_files", music_folder)
        folder_to_keep = "final"
        for folder_name in os.listdir(main_directory):
            folder_path = os.path.join(main_directory, folder_name)
            if os.path.isdir(folder_path) and folder_name != folder_to_keep:
                shutil.rmtree(folder_path)
    return (
        f"Audio file {os.path.basename(input_audio_path).split('.')[0]} converted with success",
        result,
    )


# ============================================================================
# TTS (Text-to-Speech) Functions
# ============================================================================

import asyncio
import edge_tts
from datetime import datetime
from typing import Tuple


# Edge TTS Voices by Language
EDGE_TTS_VOICES = {
    "English (US)": [
        "en-US-JennyNeural", "en-US-GuyNeural", "en-US-AriaNeural", "en-US-DavisNeural",
        "en-US-AmberNeural", "en-US-AnaNeural", "en-US-AshleyNeural", "en-US-BrandonNeural",
        "en-US-ChristopherNeural", "en-US-CoraNeural", "en-US-ElizabethNeural", "en-US-EricNeural",
        "en-US-JacobNeural", "en-US-MichelleNeural", "en-US-MonicaNeural", "en-US-RogerNeural",
    ],
    "English (UK)": [
        "en-GB-SoniaNeural", "en-GB-RyanNeural", "en-GB-LibbyNeural", "en-GB-AbbiNeural",
        "en-GB-AlfieNeural", "en-GB-BellaNeural", "en-GB-ElliotNeural", "en-GB-EthanNeural",
        "en-GB-HollieNeural", "en-GB-MaisieNeural", "en-GB-NoahNeural", "en-GB-OliverNeural",
        "en-GB-OliviaNeural", "en-GB-ThomasNeural",
    ],
    "Spanish": [
        "es-ES-ElviraNeural", "es-ES-AlvaroNeural", "es-MX-DaliaNeural", "es-MX-JorgeNeural",
        "es-AR-ElenaNeural", "es-AR-TomasNeural", "es-CO-SalomeNeural", "es-CO-GonzaloNeural",
    ],
    "French": [
        "fr-FR-DeniseNeural", "fr-FR-HenriNeural", "fr-CA-SylvieNeural", "fr-CA-JeanNeural", "fr-CA-AntoineNeural",
    ],
    "German": [
        "de-DE-KatjaNeural", "de-DE-ConradNeural", "de-DE-AmalaNeural", "de-DE-BerndNeural",
        "de-DE-ChristophNeural", "de-DE-ElkeNeural", "de-DE-GiselaNeural", "de-DE-KasperNeural",
        "de-DE-KillianNeural", "de-DE-KlarissaNeural", "de-DE-KlausNeural", "de-DE-LouisaNeural",
        "de-DE-MajaNeural", "de-DE-RalfNeural", "de-DE-TanjaNeural",
    ],
    "Portuguese (Brazil)": [
        "pt-BR-FranciscaNeural", "pt-BR-AntonioNeural", "pt-BR-BrendaNeural", "pt-BR-DonatoNeural",
        "pt-BR-ElzaNeural", "pt-BR-FabioNeural", "pt-BR-GiovannaNeural", "pt-BR-HumbertoNeural",
        "pt-BR-JulioNeural", "pt-BR-LeilaNeural", "pt-BR-LeticiaNeural", "pt-BR-ManuelaNeural",
        "pt-BR-NicolauNeural", "pt-BR-ValerioNeural", "pt-BR-YaraNeural",
    ],
    "Italian": [
        "it-IT-ElsaNeural", "it-IT-DiegoNeural", "it-IT-BenignoNeural", "it-IT-CalimeroNeural",
        "it-IT-CataldoNeural", "it-IT-FabiolaNeural", "it-IT-FiammaNeural", "it-IT-GianniNeural",
        "it-IT-ImmaNeural", "it-IT-IrmaNeural", "it-IT-LisandroNeural", "it-IT-PalmiraNeural",
        "it-IT-PierinaNeural", "it-IT-RinaldoNeural",
    ],
    "Japanese": [
        "ja-JP-NanamiNeural", "ja-JP-KeitaNeural", "ja-JP-AoiNeural", "ja-JP-DaichiNeural",
        "ja-JP-MayuNeural", "ja-JP-NaokiNeural", "ja-JP-ShioriNeural",
    ],
    "Korean": [
        "ko-KR-SunHiNeural", "ko-KR-InJoonNeural", "ko-KR-BongJinNeural", "ko-KR-GookMinNeural",
        "ko-KR-JiMinNeural", "ko-KR-SeoHyeonNeural", "ko-KR-SoonBokNeural", "ko-KR-YuJinNeural",
    ],
    "Chinese (Mandarin)": [
        "zh-CN-XiaoxiaoNeural", "zh-CN-YunxiNeural", "zh-CN-YunyangNeural", "zh-CN-XiaoyiNeural",
        "zh-CN-YunjianNeural", "zh-CN-XiaochenNeural", "zh-CN-XiaohanNeural", "zh-CN-XiaomengNeural",
        "zh-CN-XiaomoNeural", "zh-CN-XiaoqiuNeural", "zh-CN-XiaoruiNeural", "zh-CN-XiaoshuangNeural",
        "zh-CN-XiaoxuanNeural", "zh-CN-XiaoyanNeural", "zh-CN-XiaoyouNeural", "zh-CN-XiaozhenNeural",
        "zh-CN-YunfengNeural", "zh-CN-YunhaoNeural", "zh-CN-YunxiaNeural", "zh-CN-YunyeNeural", "zh-CN-YunzeNeural",
    ],
    "Russian": ["ru-RU-SvetlanaNeural", "ru-RU-DmitryNeural"],
    "Arabic": ["ar-SA-ZariyahNeural", "ar-SA-HamedNeural"],
    "Hindi": ["hi-IN-SwaraNeural", "hi-IN-MadhurNeural"],
}


# Speech rate options
TTS_RATE_OPTIONS = {
    "Very Slow (-30%)": -30,
    "Slow (-20%)": -20,
    "Slightly Slow (-10%)": -10,
    "Normal (0%)": 0,
    "Slightly Fast (+10%)": 10,
    "Fast (+20%)": 20,
    "Very Fast (+30%)": 30,
}


async def _run_edge_tts_async(
    text: str,
    voice: str,
    rate: int,
    output_path: str
) -> Tuple[bool, str]:
    """
    Internal async function to run Edge TTS.
    
    Args:
        text: Text to synthesize
        voice: Voice ID to use
        rate: Speech rate percentage
        output_path: Path to save output
    
    Returns:
        Tuple of (success, message)
    """
    try:
        rate_str = f"+{rate}%" if rate >= 0 else f"{rate}%"
        communicate = edge_tts.Communicate(text, voice, rate=rate_str)
        await communicate.save(output_path)
        return True, f"TTS generated successfully with voice: {voice}"
    except Exception as e:
        return False, f"TTS generation failed: {str(e)}"


def run_edge_tts(
    text: str,
    voice: str,
    rate: int,
    output_path: str
) -> Tuple[bool, str]:
    """
    Run Edge TTS to generate speech from text.
    
    Args:
        text: Text to convert to speech
        voice: Voice ID to use
        rate: Speech rate percentage
        output_path: Path to save output audio
    
    Returns:
        Tuple of (success, message)
    """
    return asyncio.run(_run_edge_tts_async(text, voice, rate, output_path))


def run_rvc_voice_conversion_on_tts(
    input_audio_path: str,
    output_path: str,
    model_path: str,
    index_path: str,
    pitch: int,
    pitch_extract: str,
    filter_radius: int,
    index_rate: float,
    rms_mix_rate: float,
    protect: float,
    embedder_model: str,
    devices: str,
    export_format: str,
) -> Tuple[str, str]:
    """
    Run RVC voice conversion on TTS audio.
    
    Args:
        input_audio_path: Path to TTS audio
        output_path: Path to save converted audio
        model_path: Path to RVC model
        index_path: Path to index file
        pitch: Pitch adjustment
        pitch_extract: Pitch extraction method
        filter_radius: Filter radius
        index_rate: Index rate
        rms_mix_rate: RMS mix rate
        protect: Protect value
        embedder_model: Embedder model
        devices: Device selection
        export_format: Export format
    
    Returns:
        Tuple of (status message, output path)
    """
    try:
        # Determine device
        if torch.cuda.is_available() and devices != "cpu":
            device_info = "cuda"
        else:
            device_info = "cpu"
        
        # Run full inference
        result = full_inference_program(
            model_path=model_path,
            index_path=index_path if index_path else "",
            input_audio_path=input_audio_path,
            output_path=output_path,
            export_format_rvc=export_format,
            split_audio=False,
            autotune=False,
            vocal_model="Mel-Roformer by KimberleyJSN",
            karaoke_model="Mel-Roformer Karaoke by aufr33 and viperx",
            dereverb_model="UVR-Deecho-Dereverb",
            deecho=False,
            deecho_model="UVR-Deecho-Normal",
            denoise=False,
            denoise_model="Mel-Roformer Denoise Normal by aufr33",
            reverb=False,
            vocals_volume=0,
            instrumentals_volume=0,
            backing_vocals_volume=0,
            export_format_final=export_format,
            devices=devices,
            pitch=pitch,
            filter_radius=filter_radius,
            index_rate=index_rate,
            rms_mix_rate=rms_mix_rate,
            protect=protect,
            pitch_extract=pitch_extract,
            hop_lenght=64,
            reverb_room_size=0.5,
            reverb_damping=0.5,
            reverb_wet_gain=0.33,
            reverb_dry_gain=0.4,
            reverb_width=1.0,
            embedder_model=embedder_model,
            delete_audios=False,
            use_tta=False,
            batch_size=1,
            infer_backing_vocals=False,
            infer_backing_vocals_model="",
            infer_backing_vocals_index="",
            change_inst_pitch=0,
            pitch_back=0,
            filter_radius_back=3,
            index_rate_back=0.75,
            rms_mix_rate_back=0.25,
            protect_back=0.33,
            pitch_extract_back="rmvpe",
            hop_length_back=64,
            export_format_rvc_back="WAV",
            split_audio_back=False,
            autotune_back=False,
            embedder_model_back="contentvec",
        )
        
        return result
    except Exception as e:
        logger.error(f"RVC conversion error: {e}")
        return f"RVC conversion failed: {str(e)}", ""


def run_tts_inference(
    text: str,
    language: str,
    voice: str,
    rate: int,
    use_rvc: bool,
    model_path: str,
    index_path: str,
    pitch: int,
    pitch_extract: str,
    filter_radius: int,
    index_rate: float,
    rms_mix_rate: float,
    protect: float,
    embedder_model: str,
    devices: str,
    export_format: str,
) -> Tuple[str, Optional[str], Optional[str]]:
    """
    Run complete TTS inference pipeline.
    
    Args:
        text: Text to synthesize
        language: Language category
        voice: Voice ID
        rate: Speech rate
        use_rvc: Whether to apply RVC voice conversion
        model_path: RVC model path
        index_path: RVC index path
        pitch: Pitch adjustment
        pitch_extract: Pitch extraction method
        filter_radius: Filter radius
        index_rate: Index rate
        rms_mix_rate: RMS mix rate
        protect: Protect value
        embedder_model: Embedder model
        devices: Device selection
        export_format: Export format
    
    Returns:
        Tuple of (status, tts_audio_path, final_audio_path)
    """
    # Validate input
    if not text or not text.strip():
        return "Error: Text is empty", None, None
    
    if len(text) > 5000:
        return "Error: Text too long (max 5000 characters)", None, None
    
    # Create output directory
    output_dir = os.path.join(now_dir, "audio_files", "tts_output")
    os.makedirs(output_dir, exist_ok=True)
    
    # Generate filenames
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    tts_output_path = os.path.join(output_dir, f"tts_{timestamp}.wav")
    rvc_output_path = os.path.join(output_dir, f"tts_rvc_{timestamp}.{export_format.lower()}")
    
    # Step 1: Generate TTS
    success, tts_message = run_edge_tts(text, voice, rate, tts_output_path)
    
    if not success:
        return tts_message, None, None
    
    # If RVC not enabled, return TTS output
    if not use_rvc:
        return f"TTS completed: {voice}", tts_output_path, None
    
    # Step 2: Apply RVC voice conversion
    if not model_path or not os.path.exists(model_path):
        return "RVC model not found", tts_output_path, None
    
    rvc_status, rvc_output = run_rvc_voice_conversion_on_tts(
        input_audio_path=tts_output_path,
        output_path=rvc_output_path,
        model_path=model_path,
        index_path=index_path,
        pitch=pitch,
        pitch_extract=pitch_extract,
        filter_radius=filter_radius,
        index_rate=index_rate,
        rms_mix_rate=rms_mix_rate,
        protect=protect,
        embedder_model=embedder_model,
        devices=devices,
        export_format=export_format,
    )
    
    return rvc_status, tts_output_path, rvc_output


def get_tts_voices(language: str) -> list:
    """
    Get list of available TTS voices for a language.
    
    Args:
        language: Language category
    
    Returns:
        List of voice IDs
    """
    return EDGE_TTS_VOICES.get(language, EDGE_TTS_VOICES["English (US)"])


def get_tts_languages() -> list:
    """
    Get list of available TTS languages.
    
    Returns:
        List of language names
    """
    return list(EDGE_TTS_VOICES.keys())


def get_tts_rate_options() -> dict:
    """
    Get dictionary of TTS rate options.
    
    Returns:
        Dictionary mapping rate names to values
    """
    return TTS_RATE_OPTIONS