"""
To run locally, do python app.py
"""
import os
os.environ["NUMBA_DISABLE_CUDA"] = "1"
os.environ['OMP_NUM_THREADS'] = '2'
import gradio as gr
import spaces
from huggingface_hub import hf_hub_download, login
from nemo.collections.tts.modules.magpietts_inference.utils import ModelLoadConfig, load_magpie_model

login(token=os.environ.get("HF_TOKEN"))
print("login successful")

server_port = os.getenv("SERVER_PORT")
if isinstance(server_port, str):
    server_port = int(server_port)
print(f"server_port {server_port}")

CODEC_MODEL_PATH = "nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps"

CHECKPOINT_PATH = hf_hub_download(
    repo_id="nvidia/magpie_tts_multilingual_357m",
    filename="magpie_tts_multilingual_357m.nemo",
    token=os.environ.get("HF_TOKEN"),
)
print(f"checkpoint downloaded to {CHECKPOINT_PATH}")

SPEAKER_MAP = {
    "Aria": 0,
    "Jason": 1,
    "John": 2,
    "Leo": 3,
    "Sofia": 4,
}

LANGUAGE_MAP = {
    # "Arabic-UAE": "ar-AE",
    # "Arabic-Saudi Arabia": "ar-SA",
    "Arabic-MSA": "ar-MSA",
    "Chinese": "zh",
    "English": "en",
    "French": "fr",
    "German": "de",
    "Hindi": "hi",
    "Italian": "it",
    "Japanese": "ja",
    "Korean": "ko",
    "Brazilian Portuguese": "pt-BR",
    "Spanish": "es",
    "Vietnamese": "vi",
}

BODY = """
**MagpieTTS** is NVIDIA's state-of-the-art multilingual text-to-speech system built on the [NeMo Speech](https://github.com/NVIDIA-NeMo/Speech) framework. \
It employs a two-stage pipeline architecture: a language model generates discrete acoustic tokens from text, which are then decoded into \
high-fidelity audio using a neural audio codec: [NanoCodec](https://huggingface.co/nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps). \
The model synthesizes speech in 5 different English speakers - Sofia, Aria, Jason, Leo, \
[John Van Stan](https://librivox.org/reader/9017?primary_key=9017&search_category=reader&search_page=1&search_form=get_results&search_order=alpha) \
across 12 languages: Arabic, Chinese, English, French, German, Hindi, Italian, Japanese, Korean, Portuguese, Spanish, and Vietnamese. \
The model predicts discrete audio codec tokens autoregressively using a transformer encoder-decoder architecture. It employs multi-codebook \
prediction (typically 8 codebooks) with frame stacking (factor = 2) and a local transformer for fine-grained refinement of \
high-fidelity audio generation, and leverages techniques like attention priors, classifier-free guidance (CFG), and Group Relative Policy \
Optimization (GRPO) for improved alignment.

### Key Features of the model
- **Multilingual Support** — Synthesizes natural speech in 12 languages: Arabic, Chinese, English, French, German, Hindi, Italian, Japanese, Korean, Portuguese, Spanish, and Vietnamese.
- **Expressive Voices** — Multiple voice options with emotional tones and gender variations including 4 proprietary voices and 1 public voice.
- **Text Normalization** — Built-in text normalization for handling numbers, abbreviations, and special characters for all languages.
### Resources
- 🤗 **MagpieTTS Weights**: [nvidia/magpie_tts_multilingual_357m](https://huggingface.co/nvidia/magpie_tts_multilingual_357m)
- 🤗 **NanoCodec Weights**: [nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps](https://huggingface.co/nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps)
- 💻 **NeMo Framework**: [github.com/NVIDIA-NeMo/NeMo](https://github.com/NVIDIA-NeMo/NeMo)

Note about the demo:
- Text normalization takes time to load, so if you want faster generation select "Do not apply TN."
- Text normalization works for all supported languages.
- Text normalization is required for numbers to be processed and spoken.
- When the model runs on ZeroGPU Hardware, expect slower generation because the model checkpoint is loaded everytime.
- As the model's speakers are all native English speakers, expect accented speech in the other languages.
- The current model can generate long-form speech (more than 20 seconds) in English. However, longer generations can cause timeout due to Huggingface timeout limit.
- Loan words are not supported at this time. For example, English characters in Mandarin will lead to unexpected results.
- To add custom phone pronunciations for supported languages, replace the word with it's IPA characters surrounded by '|' characters and add a space inbetween each IPA phone. For example, "Hello world from NeMo Text to Speech." could be written as "Hello world from | ˈ n ɛ m o ʊ | Text to Speech." replacing "NeMo" with "| ˈ n ɛ m o ʊ |".
  - Text normalization does not work with custom pronunciations. To mix the two, normalize the input first and then replace the word with the custom pronunciation.
- While the checkpoint was trained with 4 Arabic dialects, only ar-MSA is recommended for use.
- For the enterprise offering, see the [MagpieTTS NIM](https://build.nvidia.com/nvidia/magpie-tts-multilingual) which includes additional native voices in the supported languages, emotional speech capabilities, and optimized batch and latency inference pipeline.
"""

model_config = ModelLoadConfig(
    nemo_file=CHECKPOINT_PATH,
    codecmodel_path=CODEC_MODEL_PATH,
    legacy_codebooks=False,
    legacy_text_conditioning=False,
    hparams_from_wandb=None,
)
model = None
try:
    print(f"model_config {model_config}")
    model, _ = load_magpie_model(model_config)
    model.eval().cuda()
    print("Model loaded and ready for inference.")
except Exception as e:
    print(f"Exception in model loading: {e}")


@spaces.GPU
def demo_tts(input_text, language, speaker, apply_TN):
    global model

    if not input_text or input_text.strip() == "":
        raise gr.Error("Please enter some text to synthesize.")

    # Reload model inside GPU context if it failed at startup (ZeroGPU fallback)
    if model is None:
        try:
            model, _ = load_magpie_model(model_config)
            model.eval().cuda()
            print("Model loaded and ready for inference.")
        except Exception as e:
            raise gr.Error(f"Model failed to load: {e}")

    if speaker not in SPEAKER_MAP:
        raise gr.Error(f"Unknown speaker: {speaker!r}. Choose from {list(SPEAKER_MAP)}.")
    if language not in LANGUAGE_MAP:
        raise gr.Error(f"Unknown language: {language!r}. Choose from {list(LANGUAGE_MAP)}.")

    apply_tn_flag = apply_TN == "Apply TN"
    speaker_idx = SPEAKER_MAP[speaker]
    language_code = LANGUAGE_MAP[language]
    input_text = input_text.strip()

    if not (input_text.endswith(".") or input_text.endswith("?") or input_text.endswith("!")):
        input_text = f"{input_text}."

    try:
        import traceback
        audio, audio_len = model.do_tts(
            input_text, language=language_code, apply_TN=apply_tn_flag, speaker_index=speaker_idx
        )
        print(f"audio {audio.shape} audio_len {audio_len}")
    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"Speech synthesis failed: {e}")

    audio_np = audio[0, :audio_len[0]].cpu().numpy()

    # Return only picklable objects (int and numpy array)
    return model.sample_rate, audio_np


demo = gr.Interface(
    fn=demo_tts,
    inputs=[
        gr.Textbox(label="Text to synthesize"),
        gr.Dropdown(
            choices=LANGUAGE_MAP.keys(),
            value="English",
            label="Target Language",
            info="Select the target language for the speech to be synthesized in",
        ),
        gr.Dropdown(
            choices=["Aria", "Jason", "John", "Leo", "Sofia"],
            value="Aria",
            label="Target Speaker",
            info="Select the target speaker whose voice you would like the the speech to be synthesized in",
        ),
        gr.Radio(
            ["Apply TN", "Do not apply TN"],
            value="Do not apply TN",
            label="Apply Text Normalization",
            info="Select if you want to apply text normalization to the input text",
        ),
    ],
    outputs="audio",
    description=BODY,
    title="Text to Speech MagpieTTS Multilingual Demo",
)

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