# pip install transformers torch gradio scipy accelerate
import numpy as np
import gradio as gr
from transformers import pipeline

# 1) Text -> Audio (TTS) model
tts = pipeline("text-to-speech", model="facebook/mms-tts-eng")

# 2) Small chat model for the "conversation" part (replies in text)
chat = pipeline(
    "text-generation",
    model="Qwen/Qwen2.5-0.5B-Instruct",
    device_map="auto",
)


def speak(text):
    """Plain text -> speech."""
    out = tts(text)
    audio = np.squeeze(out["audio"])
    return out["sampling_rate"], audio


def converse(user_message, history):
    """User text -> LLM reply -> speech of the reply."""
    history = history or []
    messages = [{"role": "system", "content": "You are a friendly assistant. Keep replies short."}]
    for u, a in history:
        messages += [{"role": "user", "content": u}, {"role": "assistant", "content": a}]
    messages.append({"role": "user", "content": user_message})

    reply = chat(messages, max_new_tokens=150, do_sample=True, temperature=0.7)[0]["generated_text"][-1]["content"]
    history.append((user_message, reply))
    return history, speak(reply), ""


with gr.Blocks(title="Text to Audio") as demo:
    gr.Markdown("# Text → Audio")

    with gr.Tab("Text to Speech"):
        txt = gr.Textbox(label="Enter text", lines=3)
        btn = gr.Button("Generate audio")
        out_audio = gr.Audio(label="Output")
        btn.click(speak, txt, out_audio)

    with gr.Tab("Voice Conversation"):
        chatbot = gr.Chatbot(label="Conversation")
        msg = gr.Textbox(label="Your message")
        reply_audio = gr.Audio(label="Spoken reply", autoplay=True)
        msg.submit(converse, [msg, chatbot], [chatbot, reply_audio, msg])

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