import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from TTS.api import TTS
import gradio as gr
import tempfile

# 1. 加载 HuggingFace 中文开源语言模型
checkpoint = "Qwen/Qwen1.5-0.5B-Chat"
tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True).eval()

# 2. 初始化 TTS（支持中文）
tts = TTS(model_name="tts_models/zh-CN/baker/glow-tts", progress_bar=False, gpu=False)

def generate_story_and_audio(prompt):
    # 拼接提示词
    input_text = f"请根据以下提示生成一个简短的儿童故事（100-200字）：{prompt}"

    inputs = tokenizer(input_text, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(**inputs, max_new_tokens=200, do_sample=True)
    story = tokenizer.decode(output[0], skip_special_tokens=True)

    # 只提取模型新生成的部分
    story_clean = story.replace(input_text, "").strip()

    # 合成语音并保存为 mp3
    tmp_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    tts.tts_to_file(text=story_clean, file_path=tmp_file.name)

    return story_clean, tmp_file.name

# Gradio 界面
demo = gr.Interface(
    fn=generate_story_and_audio,
    inputs=gr.Textbox(label="请输入故事提示", placeholder="例如：一只猫在太空旅行"),
    outputs=[
        gr.Textbox(label="生成的故事"),
        gr.Audio(label="朗读音频", type="filepath")
    ],
    title="📚 开源故事生成 + 播音助手",
    description="使用 Qwen 模型生成中文故事，使用 TTS 自动朗读"
)

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