import gradio as gr
import torch
from DPTNet_eval.DPTNet_quant_sep import load_dpt_model, dpt_sep_process
import os
import soundfile as sf
import numpy as np
import librosa
import warnings

# 加載模型
model = load_dpt_model()

def separate_audio(input_wav):
    """
    Gradio Audio(filepath) → 處理 → 回傳兩個分離後的音檔路徑
    """
    file_extension = os.path.splitext(input_wav)[1].lower()
    
    # 如果是 MP3 或其他格式，先轉成 WAV
    if file_extension != ".wav":
        data, sr = sf.read(input_wav)

        # 轉單聲道
        if len(data.shape) > 1:
            data = data.mean(axis=1)

        # 重採樣到 16kHz
        if sr != 16000:
            data = librosa.resample(data, orig_sr=sr, target_sr=16000)

        # 存成 WAV
        sf.write("input.wav", data, 16000)
        wav_path = "input.wav"
    else:
        wav_path = input_wav

    # 分離語音
    outfilename = "output.wav"
    dpt_sep_process(wav_path, model=model, outfilename=outfilename)

    return (
        outfilename.replace('.wav', '_sep1.wav'),
        outfilename.replace('.wav', '_sep2.wav')
    )

# 🎯 你提供的 description 內容（已轉為 HTML）
description_html = """
<h1 align='center'><a href='https://www.twman.org/AI/ASR/SpeechSeparation' target='_blank'>中文語者分離(分割)</a></h1>
<p align='center'><b>上傳一段混音音檔，自動分離出兩個人的聲音</b></p>

<div align='center'>
  <a href='https://www.twman.org' target='_blank'>TonTon Huang Ph.D.</a> |
  <a href='https://blog.twman.org/p/deeplearning101.html' target='_blank'>手把手帶你一起踩AI坑</a> |
  <a href='https://github.com/Deep-Learning-101' target='_blank'>GitHub</a> |
  <a href='http://deeplearning101.twman.org' target='_blank'>Deep Learning 101</a>
</div>

<br>

### 🔍 使用方式：
- 上傳一段包含兩人對話的混音音檔（支援 `.mp3`, `.wav`）
- 點擊「Separate」按鈕
- 分離出兩個說話人的音軌

<br>

### 📘 相關技術文章：
<ul>
  <li><a href='https://blog.twman.org/2025/03/AIAgent.html' target='_blank'>避開 AI Agent 開發陷阱：常見問題、挑戰與解決方案 (那些 AI Agent 實戰踩過的坑)</a>：探討多種 AI Agent 工具的應用經驗與挑戰</li>
  <li><a href='https://blog.twman.org/2024/08/LLM.html' target='_blank'>白話文手把手帶你科普 GenAI</a>：淺顯介紹生成式人工智慧核心概念</li>
  <li><a href='https://blog.twman.org/2024/09/LLM.html' target='_blank'>大型語言模型直接就打完收工？</a>：回顧 LLM 領域探索歷程</li>
  <li><a href='https://blog.twman.org/2024/07/RAG.html' target='_blank'>檢索增強生成 (Retrieval-Augmented Generation, RAG) 不是萬靈丹之優化挑戰技巧</a>：探討 RAG 技術應用與挑戰</li>
  <li><a href='https://blog.twman.org/2024/02/LLM.html' target='_blank'>大型語言模型 (LLM) 入門完整指南：原理、應用與未來</a>：探討多種 LLM 工具的應用與挑戰</li>
  <li><a href='https://blog.twman.org/2023/04/GPT.html' target='_blank'>什麼是大語言模型，它是什麼？想要嗎？(Large Language Model，LLM)</a>：探討 LLM 的發展與應用</li>
  <li><a href='https://blog.twman.org/2024/11/diffusion.html' target='_blank'>ComfyUI + Stable Diffuision</a>：深入探討影像生成與分割技術的應用</li>
  <li><a href='https://blog.twman.org/2024/02/asr-tts.html' target='_blank'>ASR/TTS 開發避坑指南：語音辨識與合成的常見挑戰與對策</a>：探討 ASR 和 TTS 技術應用中的問題</li>
  <li><a href='https://blog.twman.org/2021/04/NLP.html' target='_blank'>那些自然語言處理 (NLP) 踩的坑</a>：分享 NLP 領域的實踐經驗</li>
  <li><a href='https://blog.twman.org/2021/04/ASR.html' target='_blank'>那些語音處理 (Speech Processing) 踩的坑</a>：分享語音處理領域的實務經驗</li>
  <li><a href='https://blog.twman.org/2023/07/wsl.html' target='_blank'>用PPOCRLabel來幫PaddleOCR做OCR的微調和標註</a></li>
  <li><a href='https://blog.twman.org/2023/07/HugIE.html' target='_blank'>基於機器閱讀理解和指令微調的統一信息抽取框架之診斷書醫囑資訊擷取分析</a></li>
  <li><a href='https://github.com/shibing624/pycorrector' target='_blank'>Masked Language Model (MLM) as correction BERT</a></li>
</ul>

<br><br>

📢 *本模型基於 PyTorch + Hugging Face Hub 私有模型部署*
"""

if __name__ == "__main__":
    interface = gr.Interface(
        fn=separate_audio,
        inputs=gr.Audio(type="filepath", label="請上傳混音音檔 (.mp3/.wav)"),
        outputs=[
            gr.Audio(label="語音 1"),
            gr.Audio(label="語音 2")
        ],
        title="🎙️ 語音分離 Demo - Deep Learning 101",
        description=description_html,
        allow_flagging="never"
    )

    interface.launch(debug=True)