"""
MOSS-TTS-Realtime — Hugging Face Space
Adapted from the official app.py with proper streaming via
MossTTSRealtimeStreamingSession + AudioStreamDecoder + Web Audio API.
"""

import base64
import functools
import importlib.util
import json
import os
import subprocess
import sys
import time
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Iterator, Sequence

# ── spaces (ZeroGPU — must come before other imports) ─────────────────────────
try:
    import spaces
except ImportError:
    class _SpacesFallback:
        @staticmethod
        def GPU(*_args, **_kwargs):
            def _decorator(func):
                return func
            return _decorator
    spaces = _SpacesFallback()

# ── Repo setup ─────────────────────────────────────────────────────────────────
REPO_DIR     = Path("/home/user/moss-tts")
REALTIME_DIR = REPO_DIR / "moss_tts_realtime"

def _run(*cmd, **kw):
    subprocess.check_call(list(cmd), **kw)

if not REPO_DIR.exists():
    print("[INFO] Cloning MOSS-TTS repository…")
    _run("git", "clone", "--depth=1",
         "https://github.com/OpenMOSS/MOSS-TTS.git", str(REPO_DIR))
    print("[INFO] Clone complete.")

for p in [str(REALTIME_DIR), str(REPO_DIR)]:
    if p not in sys.path:
        sys.path.insert(0, p)

# ── Core imports ───────────────────────────────────────────────────────────────
import gradio as gr
import numpy as np
import torch
import torch._dynamo
torch._dynamo.config.cache_size_limit = 64
torch._dynamo.config.suppress_errors  = True
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")  # disable compile globally — avoids torch.device ConstantVariable bug
import torchaudio
from transformers import AutoTokenizer, AutoModel
from mossttsrealtime import MossTTSRealtime, MossTTSRealtimeProcessor
from mossttsrealtime.streaming_mossttsrealtime import (
    AudioStreamDecoder,
    MossTTSRealtimeInference,
    MossTTSRealtimeStreamingSession,
)

torch.backends.cuda.enable_cudnn_sdp(False)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(True)

# ── Constants ──────────────────────────────────────────────────────────────────
MODEL_PATH        = "OpenMOSS-Team/MOSS-TTS-Realtime"
TOKENIZER_PATH    = "OpenMOSS-Team/MOSS-TTS-Realtime"
CODEC_MODEL_PATH  = "OpenMOSS-Team/MOSS-Audio-Tokenizer"
SAMPLE_RATE       = 24_000

SUPPORTED_LANGUAGES = (
    "Chinese, English, Turkish, Korean, Japanese, German, French, "
    "Spanish, Portuguese, Italian, Russian, Arabic, Hindi, "
    "Indonesian, Vietnamese, Dutch, Polish, Swedish, Danish, Norwegian"
)

# ── Dataclasses (mirrors official app.py) ─────────────────────────────────────
@dataclass(frozen=True)
class GenerationConfig:
    temperature: float
    top_p: float
    top_k: int
    repetition_penalty: float
    repetition_window: int
    do_sample: bool
    max_length: int

@dataclass(frozen=True)
class StreamingConfig:
    text_chunk_tokens: int
    decode_chunk_frames: int
    decode_overlap_frames: int
    chunk_duration: float

@dataclass(frozen=True)
class StreamEvent:
    message: str
    audio: "tuple[int, np.ndarray] | None" = None

# ── Backend loader ─────────────────────────────────────────────────────────────
@functools.lru_cache(maxsize=1)
def _load_backend(model_path: str, tokenizer_path: str, codec_path: str,
                  device_str: str, attn_impl: str):
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required.")
    device = torch.device(device_str)
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    processor = MossTTSRealtimeProcessor(tokenizer)
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    model = MossTTSRealtime.from_pretrained(
        model_path, attn_implementation=attn_impl, torch_dtype=dtype
    ).to(device).eval()
    # Disable torch.compile on the local transformer — torch.device causes
    # a ConstantVariable assertion error in torch 2.9.1+cu128 dynamo.
    
    codec = AutoModel.from_pretrained(
        codec_path, trust_remote_code=True
    ).eval().to(device)
    print("[INFO] Backend loaded.")
    return model, tokenizer, processor, codec, device

def resolve_attn(device_str: str) -> str:
    if not torch.cuda.is_available():
        return "eager"
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    if (importlib.util.find_spec("flash_attn") is not None
            and dtype in {torch.float16, torch.bfloat16}):
        major, _ = torch.cuda.get_device_capability()
        if major >= 8:
            return "flash_attention_2"
    return "sdpa"

# ── Audio helpers ──────────────────────────────────────────────────────────────
_audio_token_cache: OrderedDict = OrderedDict()
_CACHE_SIZE = 8

def _load_audio(path: Path) -> torch.Tensor:
    wav, sr = torchaudio.load(path)
    if sr != SAMPLE_RATE:
        wav = torchaudio.functional.resample(wav, sr, SAMPLE_RATE)
    if wav.shape[0] > 1:
        wav = wav.mean(dim=0, keepdim=True)
    return wav

def _extract_codes(encode_result):
    if isinstance(encode_result, dict):
        codes = encode_result["audio_codes"]
    elif isinstance(encode_result, (list, tuple)) and encode_result:
        codes = encode_result[0]
    else:
        codes = encode_result
    if isinstance(codes, np.ndarray):
        codes = torch.from_numpy(codes)
    if isinstance(codes, torch.Tensor) and codes.dim() == 3:
        codes = codes[:, 0, :] if codes.shape[1] == 1 else codes[0]
    return codes

def _encode_audio_tokens(path: Path, codec, device: torch.device,
                          chunk_duration: float) -> np.ndarray:
    key = (str(path.resolve()), int(path.stat().st_mtime_ns), chunk_duration)
    if key in _audio_token_cache:
        _audio_token_cache.move_to_end(key)
        return _audio_token_cache[key]
    with torch.inference_mode():
        wav = _load_audio(path).to(device)
        if wav.dim() == 2:
            wav = wav.unsqueeze(0)
        result = codec.encode(wav, chunk_duration=chunk_duration)
    tokens = _extract_codes(result)
    if isinstance(tokens, torch.Tensor):
        tokens = tokens.detach().cpu().numpy()
    _audio_token_cache[key] = tokens
    _audio_token_cache.move_to_end(key)
    while len(_audio_token_cache) > _CACHE_SIZE:
        _audio_token_cache.popitem(last=False)
    return tokens

def _sanitize_tokens(tokens: torch.Tensor, codebook_size: int,
                     audio_eos_token: int) -> tuple:
    if tokens.dim() == 1:
        tokens = tokens.unsqueeze(0)
    if tokens.numel() == 0:
        return tokens, False
    eos_rows    = (tokens[:, 0] == audio_eos_token).nonzero(as_tuple=False)
    invalid_rows = ((tokens < 0) | (tokens >= codebook_size)).any(dim=1)
    stop_idx = None
    if eos_rows.numel() > 0:
        stop_idx = int(eos_rows[0].item())
    if invalid_rows.any():
        inv = int(invalid_rows.nonzero(as_tuple=False)[0].item())
        stop_idx = inv if stop_idx is None else min(stop_idx, inv)
    if stop_idx is not None:
        return tokens[:stop_idx], True
    return tokens, False

def _decode_frames(audio_frames, decoder: AudioStreamDecoder,
                   codebook_size: int, audio_eos_token: int) -> Iterator[np.ndarray]:
    for frame in audio_frames:
        tokens = frame
        if tokens.dim() == 3:
            tokens = tokens[0]
        if tokens.dim() != 2:
            continue
        tokens, _ = _sanitize_tokens(tokens, codebook_size, audio_eos_token)
        if tokens.numel() == 0:
            continue
        decoder.push_tokens(tokens.detach())
        for wav in decoder.audio_chunks():
            if wav.numel() > 0:
                yield wav.detach().cpu().numpy().reshape(-1)

def _flush_decoder(decoder: AudioStreamDecoder) -> Iterator[np.ndarray]:
    final = decoder.flush()
    if final is not None and final.numel() > 0:
        yield final.detach().cpu().numpy().reshape(-1)

def _encode_chunk(sr: int, chunk: np.ndarray, idx: int) -> str:
    chunk = chunk.astype(np.float32).reshape(-1)
    return json.dumps({
        "sr": int(sr), "idx": int(idx),
        "data": base64.b64encode(chunk.tobytes()).decode("ascii"),
    })

# ── Core streaming pipeline ────────────────────────────────────────────────────
def _build_text_only_turn(processor, user_text: str,
                           prompt_tokens) -> np.ndarray:
    system_prompt = processor.make_ensemble(prompt_tokens)
    user_text_str = "<|im_end|>\n<|im_start|>user\n" + user_text + \
                    "<|im_end|>\n<|im_start|>assistant\n"
    user_tok = processor.tokenizer(user_text_str)["input_ids"]
    user_arr = np.full((len(user_tok), processor.channels + 1),
                       fill_value=processor.audio_channel_pad, dtype=np.int64)
    user_arr[:, 0] = np.asarray(user_tok, dtype=np.int64)
    return np.concatenate([system_prompt, user_arr], axis=0)

def run_stream(
    assistant_text: str,
    prompt_audio_path: "str | None",
    user_text: str,
    gen: GenerationConfig,
    streaming: StreamingConfig,
    attn_impl: str,
) -> Iterator[StreamEvent]:
    model, tokenizer, processor, codec, device = _load_backend(
        MODEL_PATH, TOKENIZER_PATH, CODEC_MODEL_PATH, "cuda:0", attn_impl
    )

    prompt_tokens = None
    if prompt_audio_path and Path(prompt_audio_path).exists():
        prompt_tokens = _encode_audio_tokens(
            Path(prompt_audio_path), codec, device, streaming.chunk_duration
        )

    inferencer = MossTTSRealtimeInference(
        model, tokenizer, max_length=gen.max_length
    )
    inferencer.reset_generation_state(keep_cache=False)

    session = MossTTSRealtimeStreamingSession(
        inferencer, processor,
        codec=codec,
        codec_sample_rate=SAMPLE_RATE,
        codec_encode_kwargs={"chunk_duration": streaming.chunk_duration},
        prefill_text_len=processor.delay_tokens_len,
        temperature=gen.temperature,
        top_p=gen.top_p,
        top_k=gen.top_k,
        do_sample=gen.do_sample,
        repetition_penalty=gen.repetition_penalty,
        repetition_window=gen.repetition_window,
    )

    if prompt_tokens is not None:
        session.set_voice_prompt_tokens(prompt_tokens)
    else:
        session.clear_voice_prompt()

    turn_input = _build_text_only_turn(processor, user_text, prompt_tokens)
    session.reset_turn(
        input_ids=turn_input, include_system_prompt=True, reset_cache=True
    )

    decoder = AudioStreamDecoder(
        codec,
        chunk_frames=streaming.decode_chunk_frames,
        overlap_frames=streaming.decode_overlap_frames,
        decode_kwargs={"chunk_duration": -1},
        device=device,
    )

    codebook_size   = int(getattr(codec, "codebook_size", 1024))
    audio_eos_token = int(getattr(inferencer, "audio_eos_token", 1026))
    text_tokens     = tokenizer.encode(assistant_text, add_special_tokens=False)
    if not text_tokens:
        raise RuntimeError("No tokens from assistant text.")

    chunk_size = max(1, streaming.text_chunk_tokens)
    token_chunks = [text_tokens[i:i+chunk_size]
                    for i in range(0, len(text_tokens), chunk_size)]

    with codec.streaming(batch_size=1):
        for chunk in token_chunks:
            audio_frames = session.push_text_tokens(chunk)
            yield from (
                StreamEvent(message="Streaming", audio=(SAMPLE_RATE, w))
                for w in _decode_frames(audio_frames, decoder,
                                        codebook_size, audio_eos_token)
            )

        audio_frames = session.end_text()
        yield from (
            StreamEvent(message="Finalizing", audio=(SAMPLE_RATE, w))
            for w in _decode_frames(audio_frames, decoder,
                                    codebook_size, audio_eos_token)
        )

        while True:
            audio_frames = session.drain(max_steps=1)
            if not audio_frames:
                break
            yield from (
                StreamEvent(message="Draining", audio=(SAMPLE_RATE, w))
                for w in _decode_frames(audio_frames, decoder,
                                        codebook_size, audio_eos_token)
            )
            if session.inferencer.is_finished:
                break

        yield from (
            StreamEvent(message="Flushing", audio=(SAMPLE_RATE, w))
            for w in _flush_decoder(decoder)
        )

    yield StreamEvent(message="Done")

# ── Gradio GPU wrapper ─────────────────────────────────────────────────────────
@spaces.GPU(duration=120)
def generate(
    user_text: str,
    assistant_text: str,
    prompt_audio,
    temperature: float, top_p: float, top_k: int,
    repetition_penalty: float, repetition_window: int,
    text_chunk_tokens: int,
    decode_chunk_frames: int, decode_overlap_frames: int,
    chunk_duration: float,
):
    if not assistant_text.strip():
        gr.Warning("Please enter assistant text to synthesize.")
        return

    attn_impl = resolve_attn("cuda:0")
    gen = GenerationConfig(
        temperature=temperature, top_p=top_p, top_k=int(top_k),
        repetition_penalty=repetition_penalty,
        repetition_window=int(repetition_window),
        do_sample=True, max_length=3000,
    )
    streaming = StreamingConfig(
        text_chunk_tokens=int(text_chunk_tokens),
        decode_chunk_frames=int(decode_chunk_frames),
        decode_overlap_frames=int(decode_overlap_frames),
        chunk_duration=float(chunk_duration),
    )

    started_at      = time.monotonic()
    first_chunk_at  = None
    full_chunks     = []
    chunk_index     = 0
    stream_reset    = json.dumps({"reset": True})
    yield stream_reset, None, "Starting…"

    try:
        for event in run_stream(
            assistant_text=assistant_text,
            prompt_audio_path=prompt_audio,
            user_text=user_text or "Hello!",
            gen=gen, streaming=streaming, attn_impl=attn_impl,
        ):
            if event.audio is None:
                yield gr.update(), gr.update(), event.message
                continue
            sr, chunk = event.audio
            chunk = np.asarray(chunk, dtype=np.float32).reshape(-1)
            if chunk.size == 0:
                continue
            if first_chunk_at is None:
                first_chunk_at = time.monotonic()
            full_chunks.append(chunk)
            chunk_index += 1
            ttfb_ms = (first_chunk_at - started_at) * 1000 if first_chunk_at else float("nan")
            payload = _encode_chunk(sr, chunk, chunk_index)
            yield payload, gr.update(), f"Streaming… chunks={chunk_index} | TTFB={ttfb_ms:.0f}ms"

        if full_chunks:
            full_audio   = np.concatenate(full_chunks)
            elapsed      = time.monotonic() - started_at
            audio_sec    = full_audio.size / SAMPLE_RATE
            rtf          = elapsed / audio_sec if audio_sec > 0 else float("inf")
            ttfb_ms      = (first_chunk_at - started_at) * 1000 if first_chunk_at else float("nan")
            done_msg = (
                f"Done | chunks={chunk_index} | audio={audio_sec:.2f}s | "
                f"elapsed={elapsed:.2f}s | RTF={rtf:.3f} | TTFB={ttfb_ms:.0f}ms"
                f"{'  ✅' if rtf < 1.0 else '  ⚠️'}"
            )
            yield gr.update(), (SAMPLE_RATE, full_audio), done_msg
        else:
            yield gr.update(), gr.update(), "Done | no audio chunks emitted"

    except Exception as exc:
        import traceback; traceback.print_exc()
        yield gr.update(), gr.update(), f"Error: {exc}"

# ── Gradio UI ──────────────────────────────────────────────────────────────────
STREAM_PLAYER_HTML = """
<style>
#pcm_stream {
  position: absolute !important; left: -9999px !important;
  width: 1px !important; height: 1px !important;
  opacity: 0 !important; pointer-events: none !important;
}
#pcm_stream textarea, #pcm_stream input {
  width: 1px !important; height: 1px !important; opacity: 0 !important;
}
</style>
<div id="pcm-stream-status" style="font-size:12px;color:#555;">
  Live playback via Web Audio API. Click Generate to unlock audio.
</div>
<div id="pcm-stream-meta" style="font-size:12px;color:#333;margin:6px 0;">
  Now Playing Chunk: <span id="pcm-stream-playing">-</span> &nbsp;|&nbsp;
  Last Yielded: <span id="pcm-stream-yielded">-</span>
</div>
"""

STREAM_PLAYER_JS = r"""
const elemId = "pcm_stream";
if (window.__pcm_streaming_inited__) return;
window.__pcm_streaming_inited__ = true;

let audioCtx = null, nextTime = 0, lastIdx = -1, lastValue = "",
    boundField = null, usingSetterHook = false;
const FADE_MS = 6, MIN_BUFFER_SEC = 0.25;
const statusEl  = document.getElementById("pcm-stream-status");
const playingEl = document.getElementById("pcm-stream-playing");
const yieldedEl = document.getElementById("pcm-stream-yielded");

function setStatus(m){ if(statusEl) statusEl.textContent = m; }
function setPlaying(i){ if(playingEl) playingEl.textContent = `${i}`; }
function setYielded(i){ if(yieldedEl) yieldedEl.textContent = `${i}`; }

function initAudio(sr){
  if(audioCtx && audioCtx.sampleRate !== sr){ audioCtx.close(); audioCtx = null; }
  if(!audioCtx){ audioCtx = new (window.AudioContext||window.webkitAudioContext)({sampleRate:sr}); nextTime = audioCtx.currentTime; }
  if(audioCtx.state === "suspended") audioCtx.resume();
}

function decodeBase64ToFloat32(b64){
  const bin = atob(b64), len = bin.length, bytes = new Uint8Array(len);
  for(let i=0;i<len;i++) bytes[i] = bin.charCodeAt(i);
  return new Float32Array(bytes.buffer);
}

function playChunk(samples, sr, idx){
  initAudio(sr);
  const buf = audioCtx.createBuffer(1, samples.length, sr);
  buf.copyToChannel(samples, 0);
  const src = audioCtx.createBufferSource(), gain = audioCtx.createGain();
  src.buffer = buf; src.connect(gain); gain.connect(audioCtx.destination);
  const now = audioCtx.currentTime;
  if(nextTime < now + MIN_BUFFER_SEC) nextTime = now + MIN_BUFFER_SEC;
  const t0 = Math.max(now, nextTime), t1 = t0 + buf.duration;
  const fade = Math.min(FADE_MS/1000, buf.duration/4);
  gain.gain.setValueAtTime(0, t0);
  gain.gain.linearRampToValueAtTime(1, t0+fade);
  gain.gain.setValueAtTime(1, Math.max(t0+fade, t1-fade));
  gain.gain.linearRampToValueAtTime(0, t1);
  src.start(t0); nextTime = t1;
  setPlaying(idx); setStatus(`Streaming… (chunk ${idx})`);
}

function handlePayloadObject(p){
  if(!p) return;
  if(p.reset){
    lastIdx=-1; lastValue="";
    if(audioCtx){ audioCtx.close(); audioCtx=null; }
    setPlaying("-"); setYielded("-");
    setStatus("Live playback via Web Audio API. Click Generate to unlock audio.");
    return;
  }
  const idx = p.idx ?? 0;
  if(idx <= lastIdx) return;
  lastIdx = idx;
  setYielded(idx);
  playChunk(decodeBase64ToFloat32(p.data), p.sr||24000, idx);
}

function handlePayload(text){
  if(!text) return;
  let p; try{ p = JSON.parse(text); } catch(e){ return; }
  if(Array.isArray(p)){ p.forEach(handlePayloadObject); } else { handlePayloadObject(p); }
}

function hookField(field){
  if(!field || field === boundField) return;
  boundField = field;
  const proto = field.tagName==="TEXTAREA" ? HTMLTextAreaElement.prototype : HTMLInputElement.prototype;
  const desc = Object.getOwnPropertyDescriptor(proto, "value");
  if(!desc||!desc.get||!desc.set){ usingSetterHook=false; return; }
  usingSetterHook = true;
  const nGet=desc.get, nSet=desc.set;
  Object.defineProperty(field,"value",{
    configurable:true,
    get(){ return nGet.call(field); },
    set(v){ nSet.call(field,v); if(v&&v!==lastValue){ lastValue=v; handlePayload(v); } }
  });
  const init = field.value;
  if(init&&init!==lastValue){ lastValue=init; handlePayload(init); }
}

function pollField(){
  const f = document.querySelector(`#${elemId} textarea, #${elemId} input`);
  if(!f){ boundField=null; usingSetterHook=false; }
  else if(f!==boundField){ hookField(f); }
  setTimeout(pollField, 300);
}

function pollValue(){
  if(usingSetterHook){ setTimeout(pollValue, 500); return; }
  const f = document.querySelector(`#${elemId} textarea, #${elemId} input`);
  if(!f){ setTimeout(pollValue,300); return; }
  const v = f.value;
  if(v&&v!==lastValue){ lastValue=v; handlePayload(v); }
  setTimeout(pollValue, 40);
}

function tryUnlockAudio(){
  if(!audioCtx) audioCtx = new (window.AudioContext||window.webkitAudioContext)();
  if(audioCtx.state==="suspended") audioCtx.resume();
}

document.addEventListener("click", e => { if(e.target.closest("#tts_generate")) tryUnlockAudio(); });
pollField(); pollValue();
"""

with gr.Blocks(title="MOSS-TTS-Realtime") as demo:
    gr.Markdown(f"""
    # 🎙️ MOSS-TTS-Realtime
    **Context-aware, multi-turn streaming TTS with real-time Web Audio playback.**
    **Supported languages:** {SUPPORTED_LANGUAGES}

    > 💡 Upload a 10–30s reference audio clip for voice cloning.
    > Without reference audio the model uses a random speaker voice.
    """)

    with gr.Row():
        with gr.Column(scale=3):
            user_text = gr.Textbox(
                label="User Text (conversation context, optional)",
                placeholder="e.g. Hello, how are you?", lines=2)
            assistant_text = gr.Textbox(
                label="Assistant Text (what to synthesize) ✱ required",
                placeholder="Enter the text you want to synthesize here…", lines=5)
            prompt_audio = gr.Audio(
                label="🎵 Reference Audio (optional — voice cloning)",
                type="filepath", sources=["upload", "microphone"])

            gr.HTML(STREAM_PLAYER_HTML, js_on_load=STREAM_PLAYER_JS)

            run_btn = gr.Button("▶ Generate", variant="primary",
                                elem_id="tts_generate")
            stream_data  = gr.Textbox(label="PCM Stream", elem_id="pcm_stream",
                                      interactive=False, lines=2)
            output_audio = gr.Audio(label="🔊 Final Audio", type="numpy")
            status       = gr.Textbox(label="📊 Status / Stats",
                                      lines=2, interactive=False)

        with gr.Column(scale=2):
            gr.Markdown("### ⚙️ Generation Parameters")
            temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.05,
                label="Temperature")
            top_p = gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="Top-P")
            top_k = gr.Slider(1, 100, value=30, step=1, label="Top-K")
            repetition_penalty = gr.Slider(1.0, 2.0, value=1.1, step=0.05,
                label="Repetition Penalty")
            repetition_window = gr.Slider(1, 200, value=50, step=1,
                label="Repetition Window")

            gr.Markdown("### 🔧 Streaming Options")
            text_chunk_tokens = gr.Slider(1, 64, value=12, step=1,
                label="Text Chunk Tokens",
                info="Tokens pushed to model per step — smaller = lower TTFB")
            decode_chunk_frames = gr.Slider(0, 20, value=12, step=1,
                label="Decode Chunk Frames")
            decode_overlap_frames = gr.Slider(0, 10, value=0, step=1,
                label="Decode Overlap Frames")
            chunk_duration = gr.Slider(0.0, 1.0, value=0.24, step=0.01,
                label="Codec Chunk Duration (s)")

    gr.Markdown("""
    ---
    ### ℹ️ Notes
    - **TTFB** = time until the first audio chunk arrives in the browser.
    - **RTF < 1.0** = faster than real-time generation.
    - Audio plays live in the browser via Web Audio API as chunks arrive.
    - The final waveform is also shown in the audio player below the status box.
    """)

    run_btn.click(
        fn=generate,
        inputs=[user_text, assistant_text, prompt_audio,
                temperature, top_p, top_k,
                repetition_penalty, repetition_window,
                text_chunk_tokens,
                decode_chunk_frames, decode_overlap_frames, chunk_duration],
        outputs=[stream_data, output_audio, status],
    )

if __name__ == "__main__":
    demo.queue(max_size=8, default_concurrency_limit=1).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT", "7860"))),
        ssr_mode=False,
        theme=gr.themes.Soft(),
    )