import gradio as gr
import torch
import torchaudio
import numpy as np
from ichigo_asr.demo.utils import load_model

# Hàm tải mô hình Ichigo Whisper với map_location=cpu
def init_model():
    try:
        # Chỉ định rõ ràng map_location='cpu' để tải mô hình trên CPU
        ichigo_model = load_model(
            ref="homebrewltd/ichigo-whisper:merge-medium-vi-2d-2560c-dim64.pth",
            size="merge-medium-vi-2d-2560c-dim64",
            map_location=torch.device('cpu')  # Thêm tham số này
        )
        device = "cpu"  # Chỉ sử dụng CPU
        ichigo_model.ensure_whisper(device)
        ichigo_model.to(device)
        return ichigo_model, device
    except Exception as e:
        print(f"Lỗi khi tải mô hình: {e}")
        return None, "cpu"

# Khởi tạo mô hình
ichigo_model, device = init_model()

def transcribe(audio_path):
    if ichigo_model is None:
        return "Không thể tải mô hình. Vui lòng kiểm tra logs."
    
    try:
        # Tải file âm thanh
        wav, sr = torchaudio.load(audio_path)
        # Chuyển đổi sang 16kHz nếu cần
        if sr != 16000:
            wav = torchaudio.functional.resample(wav, sr, 16000)
        
        # Chuyển đổi sang mono nếu là stereo
        if wav.shape[0] > 1:
            wav = wav.mean(dim=0, keepdim=True)
        
        # Đảm bảo dữ liệu nằm trên CPU
        wav = wav.to(device)
        
        # Thực hiện dự đoán
        transcribe_result = ichigo_model.inference(wav)
        
        # Trả về kết quả
        return transcribe_result[0].text
    except Exception as e:
        return f"Lỗi khi nhận dạng giọng nói: {str(e)}"

# Tạo giao diện Gradio
title = "Ichigo Whisper Speech Recognition Demo"
description = """
# 🍓 Ichigo Whisper Speech Recognition
Sử dụng mô hình Ichigo-whisper để nhận dạng giọng nói.
Mô hình này có hiệu suất tốt cho cả tiếng Anh và tiếng Việt!
"""

demo = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(sources=["microphone", "upload"], type="filepath"),
    outputs=gr.Textbox(label="Phiên âm"),
    title=title,
    description=description
)

# Khởi chạy ứng dụng
if __name__ == "__main__":
    demo.launch()