app.py

import importlib
import subprocess
import sys
import zipfile
import os
import time
import gradio as gr
from moviepy.editor import VideoFileClip, AudioFileClip
from pydub import AudioSegment
import shutil
import atexit

def install_package(name, pkg_type, upgrade=False, import_name=None, check_cmd=None):
    if pkg_type == "pip":
        import_name = import_name or name.replace("-", "_")
        try:
            importlib.import_module(import_name)
            if upgrade:
                print(f"⏫ Nâng cấp {name} ...")
                subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", name])
            else:
                print(f"✅ Đã có sẵn: {name}")
        except ImportError:
            print(f"⏳ Đang cài đặt: {name} ...")
            subprocess.check_call([sys.executable, "-m", "pip", "install", name])

install_package("gradio", "pip", upgrade=False)

import os
import time
import gradio as gr
from moviepy.editor import VideoFileClip, AudioFileClip
from pydub import AudioSegment
import shutil
import atexit

# ====== [WEBM] Các định dạng được coi là "audio-only" (xử lý như audio, bỏ qua video track) ======
WEBM_AS_AUDIO_EXTS = {'.webm'}

# Lưu danh sách temp directories để cleanup sau
temp_dirs = []

def cleanup_temp_files():
    """Xóa tất cả temporary files khi thoát"""
    for temp_dir in temp_dirs:
        try:
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)
        except Exception as e:
            print(f"Không thể xóa {temp_dir}: {e}")

atexit.register(cleanup_temp_files)

def webm_to_mp3(webm_path, output_dir=None):
    """
    [WEBM] Convert file .webm sang .mp3 bằng FFmpeg (extract audio track).
    Trả về đường dẫn file mp3 đã tạo.
    """
    if output_dir is None:
        output_dir = os.path.dirname(webm_path)
    os.makedirs(output_dir, exist_ok=True)
    file_basename = os.path.splitext(os.path.basename(webm_path))[0]
    mp3_path = os.path.join(output_dir, f"{file_basename}.mp3")
    cmd = [
        'ffmpeg',
        '-i', webm_path,
        '-vn',
        '-acodec', 'libmp3lame',
        '-q:a', '2',
        mp3_path,
        '-y'
    ]
    result = subprocess.run(cmd, capture_output=True)
    if result.returncode != 0:
        raise RuntimeError(
            f"FFmpeg không thể convert WebM sang MP3:\n{result.stderr.decode(errors='replace')}"
        )
    return mp3_path

def get_duration(file_path):
    """Lấy duration từ file"""
    if not file_path:
        return 100

    clip = None
    try:
        file_ext = os.path.splitext(file_path)[1].lower()

        if file_ext in WEBM_AS_AUDIO_EXTS:
            result = subprocess.run(
                ['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
                 '-of', 'default=noprint_wrappers=1:nokey=1', file_path],
                capture_output=True, text=True
            )
            duration_str = result.stdout.strip()
            return float(duration_str) if duration_str else 100

        if file_ext in ['.mp4', '.mkv', '.avi', '.mov', '.flv']:
            clip = VideoFileClip(file_path)
            duration = clip.duration
            clip.close()
        else:
            audio = AudioSegment.from_file(file_path)
            duration = len(audio) / 1000.0
        return duration
    except:
        return 100
    finally:
        if clip:
            clip.close()


def process_upload(file_path):
    """Xử lý file upload"""
    if not file_path:
        return None, None, "", 999999, 0

    clip = None
    try:
        file_ext = os.path.splitext(file_path)[1].lower()
        file_name = os.path.basename(file_path)
        file_size = os.path.getsize(file_path) / (1024 * 1024)

        yield None, None, "Đang xử lý file upload...", 999999, 0

        if file_ext in WEBM_AS_AUDIO_EXTS:
            temp_dir = os.path.join(os.getcwd(), "temp_webm_upload")
            os.makedirs(temp_dir, exist_ok=True)
            temp_dirs.append(temp_dir)

            mp3_path = webm_to_mp3(file_path, output_dir=temp_dir)
            duration = get_duration(file_path)
            file_type = "Audio (WebM)"
            video_file = None

            info_text = f"""📁 **Tên file:** {file_name}
🎵 **Loại:** {file_type}
💾 **Kích thước:** {file_size:.2f} MB
⏱️ **Thời lượng:** {int(duration//60)}:{int(duration%60):02d} ({duration:.1f}s)
✅ **Đã extract audio sang MP3**"""

            yield (
                mp3_path,
                video_file,
                info_text,
                duration,
                0
            )
            return

        if file_ext in ['.mp4', '.mkv', '.avi', '.mov', '.flv']:
            clip = VideoFileClip(file_path)
            duration = clip.duration
            clip.close()
            file_type = "Video"
            video_file = file_path
        elif file_ext in ['.mp3', '.wav', '.m4a', '.ogg', '.flac', '.aac']:
            audio = AudioSegment.from_file(file_path)
            duration = len(audio) / 1000.0
            file_type = "Audio"
            video_file = None
        else:
            yield None, None, f"⚠️ Định dạng file không được hỗ trợ: {file_ext}", 999999, 0
            return

        info_text = f"""📁 **Tên file:** {file_name}
🎬 **Loại:** {file_type}
💾 **Kích thước:** {file_size:.2f} MB
⏱️ **Thời lượng:** {int(duration//60)}:{int(duration%60):02d} ({duration:.1f}s)"""

        yield (
            file_path,
            video_file,
            info_text,
            duration,
            0
        )

    except Exception as e:
        yield None, None, f"❌ Lỗi: {str(e)}", 999999, 0
    finally:
        if clip:
            clip.close()

def convert_to_mp3(audio_file, video_file, bitrate, sample_rate):
    """
    Convert audio/video sang MP3 với bitrate và sample rate tùy chọn (tối ưu cho transcribe).
    """
    if not audio_file and not video_file:
        yield None, "⚠️ Không có file để convert."
        return

    try:
        input_file = video_file if video_file else audio_file
        file_basename = os.path.splitext(os.path.basename(input_file))[0]

        temp_dir = os.path.join(os.getcwd(), "temp_converted")
        os.makedirs(temp_dir, exist_ok=True)
        temp_dirs.append(temp_dir)

        mp3_file = os.path.join(temp_dir, f"{file_basename}_{bitrate}_{sample_rate}hz.mp3")

        yield None, "⏳ Đang convert sang MP3..."

        cmd = [
            'ffmpeg',
            '-i', input_file,
            '-vn',
            '-acodec', 'libmp3lame',
            '-b:a', bitrate,
            '-ar', str(sample_rate),
            mp3_file,
            '-y'
        ]
        result = subprocess.run(cmd, capture_output=True)
        if result.returncode != 0:
            yield None, f"❌ Lỗi FFmpeg: {result.stderr.decode(errors='replace')}"
            return

        file_size_kb = os.path.getsize(mp3_file) / 1024
        status_msg = (
            f"✅ Convert thành công!\n"
            f"📁 File: {os.path.basename(mp3_file)}\n"
            f"🎵 Bitrate: {bitrate} | Sample rate: {sample_rate} Hz\n"
            f"💾 Kích thước: {file_size_kb:.1f} KB"
        )
        yield mp3_file, status_msg

    except Exception as e:
        import traceback
        yield None, f"❌ Lỗi: {str(e)}\n{traceback.format_exc()}"

def clear_outputs():
    return (
        gr.update(value=None, visible=False),  # source_mp3_output
        gr.update(value=None, visible=False),  # converted_mp3_output
        gr.update(value=None, visible=False),  # cut_zip_output
        "",                                    # cut_status
        gr.update(visible=False)               # clear_output_btn
    )

def cut_file_by_duration(audio_file, minutes_per_chunk):
    """Cắt file mp3 thành các đoạn theo số phút"""
    if not audio_file:
        return [], "⚠️ Không có file audio để cắt."
    try:
        import datetime
        chunk_secs = int(minutes_per_chunk) * 60
        total_duration = get_duration(audio_file)
        if total_duration <= chunk_secs:
            return [], f"⚠️ File ngắn hơn {minutes_per_chunk} phút ({total_duration:.0f}s). Không cần cắt."
        temp_dir = os.path.join(os.getcwd(), "temp_cut")
        os.makedirs(temp_dir, exist_ok=True)
        temp_dirs.append(temp_dir)
        ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        file_basename = os.path.splitext(os.path.basename(audio_file))[0]
        output_files = []
        part = 1
        start = 0.0
        while start < total_duration:
            end = min(start + chunk_secs, total_duration)
            out_path = os.path.join(temp_dir, f"{file_basename}_part{part:02d}_{ts}.mp3")
            cmd = ['ffmpeg', '-ss', str(start), '-i', audio_file,
                   '-t', str(end - start), '-acodec', 'libmp3lame', '-q:a', '2', out_path, '-y']
            result = subprocess.run(cmd, capture_output=True)
            if result.returncode == 0:
                output_files.append(out_path)
            start += chunk_secs
            part += 1
        status = f"✅ Đã cắt thành {len(output_files)} đoạn × {minutes_per_chunk} phút"
        return output_files, status
    except Exception as e:
        import traceback
        return [], f"❌ Lỗi: {str(e)}\n{traceback.format_exc()}"


def cut_file_by_parts(audio_file, num_parts):
    """Cắt file mp3 thành N phần bằng nhau"""
    if not audio_file:
        return [], "⚠️ Không có file audio để cắt."
    try:
        import datetime
        num_parts = int(num_parts)
        total_duration = get_duration(audio_file)
        chunk_secs = total_duration / num_parts
        temp_dir = os.path.join(os.getcwd(), "temp_cut")
        os.makedirs(temp_dir, exist_ok=True)
        temp_dirs.append(temp_dir)
        ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        file_basename = os.path.splitext(os.path.basename(audio_file))[0]
        output_files = []
        for part in range(1, num_parts + 1):
            start = (part - 1) * chunk_secs
            out_path = os.path.join(temp_dir, f"{file_basename}_part{part:02d}of{num_parts}_{ts}.mp3")
            cmd = ['ffmpeg', '-ss', str(start), '-i', audio_file,
                   '-t', str(chunk_secs), '-acodec', 'libmp3lame', '-q:a', '2', out_path, '-y']
            result = subprocess.run(cmd, capture_output=True)
            if result.returncode == 0:
                output_files.append(out_path)
        status = f"✅ Đã cắt thành {len(output_files)} phần (~{chunk_secs/60:.1f} phút/phần)"
        return output_files, status
    except Exception as e:
        import traceback
        return [], f"❌ Lỗi: {str(e)}\n{traceback.format_exc()}"


def do_cut_file(audio_file, cut_mode, minutes_val, parts_val):
    """Dispatcher: cắt theo phút hoặc theo phần, đóng gói zip"""
    if not audio_file:
        yield gr.update(value=None, visible=False), "⚠️ Không có file audio."
        return
    yield gr.update(visible=False), "⏳ Đang cắt file..."
    if cut_mode == "Theo phút":
        files, status = cut_file_by_duration(audio_file, minutes_val)
    else:
        files, status = cut_file_by_parts(audio_file, parts_val)
    if not files:
        yield gr.update(value=None, visible=False), status
        return
    import zipfile, datetime
    ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    file_basename = os.path.splitext(os.path.basename(audio_file))[0]
    zip_dir = os.path.join(os.getcwd(), "temp_cut")
    os.makedirs(zip_dir, exist_ok=True)
    zip_path = os.path.join(zip_dir, f"{file_basename}_cut_{ts}.zip")
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
        for f in files:
            zf.write(f, os.path.basename(f))
    zip_size_mb = os.path.getsize(zip_path) / (1024 * 1024)
    status += f"\n📦 ZIP: {os.path.basename(zip_path)} ({zip_size_mb:.1f} MB)"
    yield gr.update(value=zip_path, visible=True), status


# ====== GRADIO INTERFACE ======

css = """
#textbox_id textarea {
  color: black !important;
  font-size: 16px !important;
  font-family: 'IBM Plex Sans', sans-serif !important;
}

#textbox_id placeholder::textarea {
  color: black !important;
  font-size: 16px !important;
  font-family: 'IBM Plex Sans', sans-serif !important;
}

#method_dropdown .menu button {
  color: *primary_50 !important;
  font-size: 16px !important;
}
"""

theme = gr.themes.Default().set(
    block_background_fill='*primary_50',
    block_border_color='*button_primary_border_color',
    block_label_text_color='*secondary_600',
    block_info_text_color='*primary_700',
    block_title_text_color='*primary_700',
    body_text_size='*text_lg'
)

with gr.Blocks(theme=theme, css=css, title="Media Transcriber with Trimming") as demo:

    gr.Markdown("# 🎤 Media Transcriber with Trimming")
    gr.Markdown("Download từ YouTube (convert MP3) hoặc upload file, cắt (tùy chọn), và transcribe với faster-whisper")

    with gr.Tab("Transcription"):
        current_file = gr.State()
        video_file_state = gr.State()

        with gr.Row():
            with gr.Column(scale=2):
                upload_file = gr.File(
                    label="Upload file audio/video từ máy tính (.mp3 .wav .m4a .ogg .flac .aac .webm .mp4 .mkv .avi .mov .flv)",
                    file_types=["video", "audio", ".webm"],
                    type="filepath"
                )
                media_info = gr.Markdown(value="Nhập URL hoặc upload file để xem thông tin...")

                with gr.Row():
                    bitrate_dropdown = gr.Dropdown(
                        choices=["32k", "64k"],
                        value="32k",
                        label="🎚️ Bitrate",
                        info="32k: nhỏ hơn | 64k: tốt hơn một chút",
                        scale=1,
                    )
                    sample_rate_dropdown = gr.Dropdown(
                        choices=["8000", "16000"],
                        value="16000",
                        label="📊 Sample Rate (Hz)",
                        info="8kHz: tối thiểu | 16kHz: khuyến nghị cho transcribe",
                        scale=1,
                    )

                convert_btn = gr.Button("🎵 Convert MP3", variant="secondary", size="lg")
                convert_status = gr.Markdown(value="")

                transcribe_btn = gr.Button("🎙️ Transcribe", variant="primary", size="lg", interactive=False)

                with gr.Accordion("🔧 Phương thức & Timestamps", open=False):
                    with gr.Row():
                        method_dropdown = gr.Dropdown(
                            choices=["model.transcribe", "BatchedInferencePipeline"],
                            label="Phương thức",
                            scale=3,
                            elem_id="method_dropdown",
                            value="BatchedInferencePipeline",
                            info="model.transcribe: chất lượng cao | Batched: nhanh hơn"
                        )
                        beam_size_dropdown = gr.Dropdown(
                            choices=["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"],
                            label="Beam Size",
                            scale=2,
                            elem_id="method_dropdown",
                            value="3",
                            info="Cao hơn = chính xác hơn nhưng chậm hơn"
                        )
                    include_timestamps_checkbox = gr.Checkbox(
                        label="Bao gồm Timestamps (HH:MM:SS)",
                        value=True,
                        info="Bao gồm timestamps trong kết quả transcription"
                    )

                with gr.Accordion("⚙️ Tham số nâng cao (Batch & VAD)", open=False):
                    with gr.Row():
                        batch_size_slider = gr.Slider(
                            minimum=1,
                            maximum=64,
                            value=16,
                            step=1,
                            label="Batch Size (BatchedInferencePipeline)",
                            info="Số lượng đoạn âm thanh xử lý cùng lúc. Ảnh hưởng đến tốc độ và bộ nhớ GPU.",
                            visible=True
                        )
                    with gr.Row():
                        min_silence_duration_ms_slider = gr.Slider(
                            minimum=0,
                            maximum=2000,
                            value=500,
                            step=50,
                            label="VAD: Min Silence Duration (ms)",
                            info="Thời lượng im lặng tối thiểu để tách phân đoạn. Ảnh hưởng đến việc phát hiện câu/từ."
                        )
                        speech_pad_ms_slider = gr.Slider(
                            minimum=0,
                            maximum=1000,
                            value=400,
                            step=50,
                            label="VAD: Speech Pad (ms)",
                            info="Thêm thời gian vào đầu/cuối mỗi phân đoạn giọng nói. Giúp giữ lại bối cảnh."
                        )
                    with gr.Row():
                        no_speech_threshold_slider = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            value=0.55,
                            step=0.05,
                            label="VAD: No Speech Threshold",
                            info="Ngưỡng xác định khi nào không có lời nói."
                        )
                        condition_on_previous_text_checkbox = gr.Checkbox(
                            label="Condition on Previous Text",
                            value=False,
                            info="Sử dụng văn bản trước đó làm điều kiện để cải thiện tính nhất quán."
                        )
                    with gr.Row():
                        minimum_speech_duration_slider = gr.Slider(
                            minimum=0.0,
                            maximum=5.0,
                            value=0.1,
                            step=0.05,
                            label="VAD: Minimum Speech Duration (s)",
                            info="Thời lượng tối thiểu của một đoạn giọng nói. Giúp lọc các âm thanh ngắn, nhiễu."
                        )

                with gr.Accordion("✂️ Cắt file MP3", open=False):
                    cut_mode_radio = gr.Radio(
                        choices=["Theo phút", "Theo phần"],
                        value="Theo phút",
                        label="Chế độ cắt",
                        info="Chỉ chọn một chế độ"
                    )
                    with gr.Row():
                        cut_minutes_dropdown = gr.Dropdown(
                            choices=["3", "5", "10", "15"],
                            value="5",
                            label="⏱️ Phút mỗi đoạn",
                            info="Áp dụng khi chọn Theo phút",
                            interactive=True,
                            scale=1,
                        )
                        cut_parts_dropdown = gr.Dropdown(
                            choices=["2", "3", "5", "10"],
                            value="2",
                            label="🔢 Số phần",
                            info="Áp dụng khi chọn Theo phần",
                            interactive=False,
                            scale=1,
                        )
                    cut_btn = gr.Button("✂️ Cut File", variant="secondary", size="lg")
                    cut_status = gr.Markdown(value="")

            with gr.Column(scale=3):
                placeholder_text = (
                "📖 Hướng dẫn: Nhập Youtube video URL và bấm ENTER \n\n"
                "📖 Hoặc: Upload file từ máy tính \n\n"
                )
                transcript_output = gr.Textbox(
                    label="Transcript",
                    lines=20,
                    max_lines=50,
                    show_copy_button=True,
                    placeholder=placeholder_text,
                    autoscroll=True,
                    elem_id="textbox_id",
                )
                statistics_output = gr.Textbox(
                    label="Thông tin transcription",
                    lines=5,
                    max_lines=20,
                    show_copy_button=True,
                    visible=True,
                )
                clear_output_btn = gr.Button("🗑️ Clear Output", variant="stop", size="sm", visible=False)

                with gr.Row():
                    source_mp3_output = gr.File(label="📥 Source MP3", visible=False)
                    converted_mp3_output = gr.File(label="📥 Converted MP3", visible=False)

                with gr.Row():
                    file_output = gr.File(label="📄 TXT", visible=False)
                    srt_output = gr.File(label="🎬 SRT", visible=False)
                    vtt_output = gr.File(label="🌐 VTT", visible=False)

                with gr.Row():
                    trimmed_mp3_output = gr.File(label="✂️ Trimmed MP3", visible=False)
                    trimmed_mp4_output = gr.File(label="✂️ Trimmed Video", visible=False)

                cut_zip_output = gr.File(label="📦 Cut ZIP", visible=False)

        # ====== EVENT HANDLERS ======

        # File upload
        upload_file.change(
            fn=process_upload,
            inputs=[upload_file],
            outputs=[current_file, video_file_state, media_info, gr.State(999999), gr.State(0)]
        )

        # Show source MP3 after file is set
        current_file.change(
            fn=lambda f: gr.update(value=f, visible=True) if f else gr.update(visible=False),
            inputs=[current_file],
            outputs=[source_mp3_output]
        )

        # ====== NÚT CONVERT MP3 ======
        convert_btn.click(
            fn=lambda: (gr.update(interactive=False), gr.update(value="⏳ Đang convert...")),
            inputs=None,
            outputs=[convert_btn, convert_status]
        ).then(
            fn=convert_to_mp3,
            inputs=[current_file, video_file_state, bitrate_dropdown, sample_rate_dropdown],
            outputs=[converted_mp3_output, convert_status]
        ).then(
            fn=lambda f: (gr.update(interactive=True), gr.update(visible=True) if f else gr.update(visible=False)),
            inputs=[converted_mp3_output],
            outputs=[convert_btn, converted_mp3_output]
        )

        # ====== CUT FILE ======
        cut_mode_radio.change(
            fn=lambda mode: (
                gr.update(interactive=(mode == "Theo phút")),
                gr.update(interactive=(mode == "Theo phần"))
            ),
            inputs=[cut_mode_radio],
            outputs=[cut_minutes_dropdown, cut_parts_dropdown]
        )

        cut_btn.click(
            fn=lambda: gr.update(interactive=False),
            inputs=None,
            outputs=[cut_btn]
        ).then(
            fn=do_cut_file,
            inputs=[current_file, cut_mode_radio, cut_minutes_dropdown, cut_parts_dropdown],
            outputs=[cut_zip_output, cut_status]
        ).then(
            fn=lambda: gr.update(interactive=True),
            inputs=None,
            outputs=[cut_btn]
        )

        # Clear Output button handler
        clear_output_btn.click(
            fn=clear_outputs,
            inputs=None,
            outputs=[
                source_mp3_output,
                converted_mp3_output,
                cut_zip_output,
                cut_status,
                clear_output_btn
            ]
        )

# ====== QUEUE + LAUNCH ======
demo.queue(default_concurrency_limit=2)
demo.launch()