app.py

import gradio as gr
from transformers import pipeline # 保持使用 transformers pipeline
import torch
import numpy as np
import plotly.graph_objects as go
import os
import scipy.io.wavfile # For saving/reading wav files

# --- 全局变量和模型加载 ---
HF_HOME = os.getenv("HF_HOME", "./hf_cache")
os.makedirs(HF_HOME, exist_ok=True)

# 定义模型路径和名称
# 模型 A: 文本到语音 (Text-to-Speech) - 更换为 SpeechBrain 的轻量级模型
MODEL_TTS = "speechbrain/tts-tacotron2-ljspeech" 

# 模型 B: 文本到音效/简单音频 (Text-to-Audio/Sound Generation) - 仍然不加载
MODEL_AUDIOGEN = "facebook/audiogen-small" 

# 确定加载模型的设备
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"当前使用的设备: {device}")

# --- 模型加载函数 ---
tts_pipeline = None
audiogen_pipeline = None # 保持声明，但不会加载

def load_tts_model():
    """加载 MMS-TTS 模型。现在尝试 SpeechBrain 的 Tacotron2。"""
    print(f"正在加载 TTS 模型: {MODEL_TTS} 到 {device}...")
    try:
        # 对于 SpeechBrain 的模型，直接使用 pipeline 可能更简单
        # 注意：SpeechBrain 的 pipeline 可能需要安装 speechbrain 库
        # 确保你的 requirements.txt 包含 speechbrain
        pipe = pipeline("text-to-speech", model=MODEL_TTS, device=device)
        print("TTS 模型加载成功。")
        return pipe
    except Exception as e:
        print(f"加载 TTS (SpeechBrain) 时出错: {e}")
        return None

# !!! 注意：这个函数仍然存在，但不会在 initialize_all_models 中调用 !!!
def load_audiogen_model():
    """加载 AudioGen-small 模型。"""
    print(f"正在加载 AudioGen 模型: {MODEL_AUDIOGEN} 到 {device}...")
    try:
        pipe = pipeline("text-to-audio", model=MODEL_AUDIOGEN, device=device)
        print("AudioGen 模型加载成功。")
        return pipe
    except Exception as e:
        print(f"加载 AudioGen 时出错: {e}")
        return None

def initialize_all_models():
    """在 Gradio 界面加载时调用，现在只初始化 TTS 模型。"""
    global tts_pipeline, audiogen_pipeline
    tts_pipeline = load_tts_model()
    # audiogen_pipeline = load_audiogen_model() # <--- !!! 这一行已被注释掉，不再加载 AudioGen !!!

# --- 音频生成推理函数 ---

def generate_audio_tts(prompt):
    """使用 TTS 模型生成语音。"""
    if tts_pipeline is None:
        print("TTS 模型未加载或加载失败。")
        return (16000, np.zeros(0)) # 返回空音频
    try:
        result = tts_pipeline(prompt)
        audio_array = result['audio']
        sampling_rate = result['sampling_rate']
        return (sampling_rate, audio_array)
    except Exception as e:
        print(f"使用 TTS 生成音频出错: {e}")
        return (16000, np.zeros(0))

def generate_audio_audiogen(prompt):
    """使用 AudioGen-small 生成音频。由于不再加载，这里会返回空音频并提示。"""
    if audiogen_pipeline is None:
        print("AudioGen 模型未加载或加载失败 (在CPU环境下可能已禁用)。")
        return (44100, np.zeros(0)) # 返回空音频，AudioGen通常用44.1kHz
    try:
        audio_output = audiogen_pipeline(prompt, sampling_rate=16000)
        audio_array = audio_output['audio'][0].cpu().numpy() if isinstance(audio_output['audio'], torch.Tensor) else audio_output['audio']
        sampling_rate = audio_output['sampling_rate'] if 'sampling_rate' in audio_output else 16000
        return (sampling_rate, audio_array)
    except Exception as e:
        print(f"使用 AudioGen 生成音频出错: {e}")
        return (16000, np.zeros(0))

# --- Arena 选项卡逻辑 ---
def arena_predict(prompt):
    """Arena 选项卡的主预测函数。只加载 TTS 模型时，AudioGen 将返回空。"""
    print(f"收到提示词: {prompt}")
        
    # TTS 生成
    tts_sr, tts_audio = generate_audio_tts(prompt)
        
    # AudioGen-small 生成 (此处会因为 audiogen_pipeline 为 None 而返回空音频)
    audiogen_sr, audiogen_audio = generate_audio_audiogen(prompt)
    
    return (tts_sr, tts_audio), (audiogen_sr, audiogen_audio)

# --- 模型对比/GRACE评估逻辑 ---
# 假设的 GRACE 评估数据
grace_data = {
    "MMS-TTS-Eng": { # 虽然现在用 SpeechBrain，但报表里仍沿用原名
        "Generalization": 2.5, 
        "Relevance": 4.5,     
        "Artistry": 4.0,      
        "Efficiency": 4.5     
    },
    "AudioGen-small": {
        "Generalization": 3.5, 
        "Relevance": 3.5,      
        "Artistry": 3.5,       
        "Efficiency": 4.0      
    }
}

def create_grace_radar_chart():
    """创建 GRACE 框架的雷达图。"""
    categories = ['Generalization', 'Relevance', 'Artistry', 'Efficiency']
        
    fig = go.Figure()

    for model_name, scores in grace_data.items():
        values = [scores[cat] for cat in categories]
        values.append(values[0]) # 闭合雷达图
                
        fig.add_trace(go.Scatterpolar(
            r=values,
            theta=categories + [categories[0]],
            fill='toself',
            name=model_name,
            line_color='blue' if model_name == "MMS-TTS-Eng" else 'red'
        ))

    fig.update_layout(
        polar=dict(
            radialaxis_tickangle=45,
            radialaxis_tickfont_size=12,
            radialaxis=dict(
                visible=True,
                range=[0, 5],
                tickvals=[1, 2, 3, 4, 5],
                ticktext=['1', '2', '3', '4', '5']
            ),
            angularaxis=dict(
                direction="clockwise",
                period=6
            )
        ),
        showlegend=True,
        title="GRACE 框架模型对比"
    )
        
    return fig

# --- 报告内容 (Markdown 格式) ---
report_content = """
# 文本到音频生成模型对比实验报告

## 1. 模型及类别选择

### 1.1 所选模型的类型及背景说明
本次实验聚焦于**文本到音频（Text-to-Audio）生成模型**，具体选择了两个在功能上有所侧重且模型尺寸相对较小的模型进行对比：
1.  **`speechbrain/tts-tacotron2-ljspeech` (新的 TTS 模型)**：这是一个由 SpeechBrain 开发的**文本到语音（Text-to-Speech, TTS）**模型，基于 Tacotron2 架构。相比之前尝试的模型，它通常更小巧，对计算资源的需求可能更低，因此更适合在有限的 CPU 环境下运行。
2.  **`facebook/audiogen-small` (AudioGen-small)**：这款模型是 Meta AI AudioCraft 框架下的一个较小版本，专注于**文本到音效或简单音频片段**的生成。它能够根据文本描述创建环境音、乐器声或短小的声音事件。

### 1.2 模型用途对比简述
* **TTS 模型 (Tacotron2)** 的主要用途是将书面文本转换为自然发音的**人类语音**，适用于有声读物、语音助手、导航系统等场景。
* **AudioGen-small** 的主要用途是根据文本描述生成各种**非语音的音效或短音乐片段**，例如“雨声”、“鸟鸣”或“简单的吉他旋律”，适用于游戏音效、视频配乐或环境音模拟。

### 1.3 两个模型异同点分析
**相同点：**
* 都属于文本到音频生成领域，将文本作为输入。
* 都可在 Hugging Face Hub 上找到。
* 都相对轻量级，相比大型文生音/乐模型对资源要求更低。

**不同点：**
* **核心功能：** TTS 模型专注于**语音合成**，目标是还原人类语音的自然度；AudioGen 专注于**非语音音频生成**，目标是模拟真实世界的音效或创作音乐片段。
* **输出类型：** TTS 模型输出的是人类语言的声音；AudioGen 输出的是环境音、乐器音、抽象音效等。
* **内部机制：** 尽管都基于 Transformer 架构（Tacotron2 包含编码器和解码器），但它们的内部训练数据、任务目标和具体架构优化不同，以适应各自的生成任务。

---

## 2. 系统实现细节
本系统通过 Gradio 构建了一个交互式 Hugging Face Space，实现了模型的统一输入和多模型输出展示。

### 2.1 Gradio 交互界面截图
（此处应放置你的 Gradio Space 的截图，显示 "Arena" 和 "模型对比" 选项卡的用户界面）
![Gradio Interface Screenshot](https://i.imgur.com/your-space-screenshot-link.png)

### 2.2 输入与输出流程图 (Mermaid 语法)
```mermaid
graph TD
    A[用户输入文本提示词] --> B{Gradio UI};
    B --> C1[TTS 模型];
    B --> C2[AudioGen-small 模型];
    C1 -- 生成语音 --> D1[TTS 音频输出];
    C2 -- 生成音效 --> D2[AudioGen-small 音频输出];
    D1 --> E[Arena 选项卡展示];
    D2 --> E;
    E -- 用户评估 --> F[GRACE 评估数据];
    F --> G[模型对比选项卡];
    G -- 生成雷达图 --> H[报告选项卡嵌入雷达图];"""