import torch
import streamlit as st
import numpy as np
import torchaudio
import torch.nn as nn
import torch.nn.functional as F
from model import SpeechSeparationModel

# ตรวจสอบว่าไฟล์โมเดลอยู่ที่ไหน
MODEL_PATH = "speech_separation_model (1).pth"

# โหลดโมเดล Speech Separation
@st.cache_resource
def load_model():
    model = SpeechSeparationModel()  # ตรวจสอบว่า SpeechSeparationModel ถูก import
    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
    model.eval()
    return model

model = load_model()

# ฟังก์ชันแยกเสียงพูดโดยใช้โมเดล
def separate_speech(audio_signal, sample_rate):
    audio_tensor = torch.tensor(audio_signal).float().unsqueeze(0).unsqueeze(0)  # (batch, channel, length)
    with torch.no_grad():
        separated_audio = model(audio_tensor).squeeze().numpy()
    return separated_audio