app.py

import gradio as gr
from transformers import MarianMTModel, MarianTokenizer, pipeline
import pysrt
import tempfile
from tqdm import tqdm
from langdetect import detect
import os
from datetime import timedelta

# Danh sách các ngôn ngữ và model tương ứng
LANGUAGE_MODELS = {
    "Tiếng Anh": "en",
    "Tiếng Việt": "vi",
    "Tiếng Pháp": "fr",
    "Tiếng Đức": "de",
    "Tiếng Tây Ban Nha": "es",
    "Tiếng Bồ Đào Nha": "pt",
    "Tiếng Ý": "it",
    "Tiếng Nga": "ru",
    "Tiếng Hà Lan": "nl",
    "Tiếng Thụy Điển": "sv",
    "Tiếng Phần Lan": "fi",
    "Tiếng Đan Mạch": "da",
    "Tiếng Na Uy": "no",
    "Tiếng Ba Lan": "pl",
    "Tiếng Séc": "cs",
    "Tiếng Hungary": "hu",
    "Tiếng Romania": "ro",
    "Tiếng Hy Lạp": "el",
    "Tiếng Thổ Nhĩ Kỳ": "tr",
    "Tiếng Hindi": "hi",
    "Tiếng Ả Rập": "ar",
    "Tiếng Trung (Giản thể)": "zh",
    "Tiếng Nhật": "ja",
    "Tiếng Hàn": "ko"
}

# Đảo ngược dictionary để lấy code từ tên ngôn ngữ
LANGUAGE_CODES = {v: k for k, v in LANGUAGE_MODELS.items()}

# Cache models để tăng tốc độ
model_cache = {}
detector = pipeline("text-classification", model="papluca/xlm-roberta-base-language-detection")

def detect_subtitle_language(file_path):
    try:
        subs = pysrt.open(file_path)
        sample_text = " ".join([sub.text for sub in subs[:10] if sub.text.strip()])
        
        if not sample_text:
            return "en"  # Mặc định là tiếng Anh nếu không phát hiện được
        
        try:
            # Sử dụng langdetect cho đơn giản
            lang_code = detect(sample_text)
            return lang_code
        except:
            # Fallback sử dụng model xlm-roberta
            result = detector(sample_text[:512])[0]  # Giới hạn độ dài đầu vào
            return result['label'].split('__')[-1]
    except Exception as e:
        print(f"Error detecting language: {e}")
        return "en"

def get_model(source_lang, target_lang):
    model_key = f"{source_lang}-{target_lang}"
    
    if model_key not in model_cache:
        model_name = f"Helsinki-NLP/opus-mt-{model_key}"
        try:
            tokenizer = MarianTokenizer.from_pretrained(model_name)
            model = MarianMTModel.from_pretrained(model_name)
            model_cache[model_key] = (model, tokenizer)
        except:
            # Fallback: Dịch qua tiếng Anh nếu không có model trực tiếp
            if source_lang != "en":
                # Dịch từ ngôn ngữ nguồn -> tiếng Anh -> ngôn ngữ đích
                model1_name = f"Helsinki-NLP/opus-mt-{source_lang}-en"
                model2_name = f"Helsinki-NLP/opus-mt-en-{target_lang}"
                
                tokenizer1 = MarianTokenizer.from_pretrained(model1_name)
                model1 = MarianMTModel.from_pretrained(model1_name)
                tokenizer2 = MarianTokenizer.from_pretrained(model2_name)
                model2 = MarianMTModel.from_pretrained(model2_name)
                
                model_cache[model_key] = ((model1, tokenizer1), (model2, tokenizer2))
            else:
                raise gr.Error(f"Không tìm thấy model dịch từ {source_lang} sang {target_lang}")
    
    return model_cache[model_key]

def translate_text(text, model, tokenizer):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    translated = model.generate(**inputs)
    return tokenizer.batch_decode(translated, skip_special_tokens=True)[0]

def add_time_to_subtitle(input_file, hours, minutes, seconds):
    if input_file is None:
        raise gr.Error("Vui lòng upload file phụ đề!")
    
    try:
        if not os.path.exists(input_file):
            raise gr.Error("File không tồn tại hoặc không thể đọc!")
        
        subs = pysrt.open(input_file)
        
        # Chuyển đổi thời gian nhập vào thành mili giây (hỗ trợ số thập phân)
        try:
            seconds_float = float(seconds)
        except ValueError:
            seconds_float = 0
            
        total_milliseconds = int((int(hours) * 3600 + int(minutes) * 60 + seconds_float) * 1000)
        
        # Thêm thời gian vào tất cả các phụ đề
        if total_milliseconds > 0:
            for sub in subs:
                sub.start.ordinal += total_milliseconds
                sub.end.ordinal += total_milliseconds
        
        # Lưu file tạm
        output_path = tempfile.NamedTemporaryFile(suffix=".srt", delete=False).name
        subs.save(output_path, encoding='utf-8')
        return output_path, f"Đã thêm {hours}h {minutes}m {seconds_float}s vào file gốc"
    
    except Exception as e:
        raise gr.Error(f"Có lỗi xảy ra khi thêm thời gian: {str(e)}")

def translate_subtitle(input_file, source_language, target_language, hours, minutes, seconds):
    if input_file is None:
        raise gr.Error("Vui lòng upload file phụ đề!")
    
    try:
        if not os.path.exists(input_file):
            raise gr.Error("File không tồn tại hoặc không thể đọc!")
        
        source_code = LANGUAGE_MODELS.get(source_language, "en")
        target_code = LANGUAGE_MODELS[target_language]
        
        model_info = get_model(source_code, target_code)
        
        subs = pysrt.open(input_file)
        
        # Chuyển đổi thời gian nhập vào thành mili giây (hỗ trợ số thập phân)
        try:
            seconds_float = float(seconds)
        except ValueError:
            seconds_float = 0
            
        total_milliseconds = int((int(hours) * 3600 + int(minutes) * 60 + seconds_float) * 1000)
        
        # Thêm thời gian vào tất cả các phụ đề
        if total_milliseconds > 0:
            for sub in subs:
                sub.start.ordinal += total_milliseconds
                sub.end.ordinal += total_milliseconds
        
        # Xử lý dịch thuật
        if isinstance(model_info[0], tuple):
            # Dịch qua tiếng Anh
            model1, tokenizer1 = model_info[0]
            model2, tokenizer2 = model_info[1]
            
            for sub in tqdm(subs, desc="Đang dịch"):
                if sub.text.strip():
                    en_text = translate_text(sub.text, model1, tokenizer1)
                    sub.text = translate_text(en_text, model2, tokenizer2)
        else:
            # Dịch trực tiếp
            model, tokenizer = model_info
            for sub in tqdm(subs, desc="Đang dịch"):
                if sub.text.strip():
                    sub.text = translate_text(sub.text, model, tokenizer)
        
        # Lưu file tạm
        output_path = tempfile.NamedTemporaryFile(suffix=".srt", delete=False).name
        subs.save(output_path, encoding='utf-8')
        return output_path, f"Dịch từ {source_language} sang {target_language} thành công! Đã thêm {hours}h {minutes}m {seconds_float}s"
    
    except Exception as e:
        raise gr.Error(f"Có lỗi xảy ra: {str(e)}")

# Giao diện Gradio
with gr.Blocks(title="Subtitle Translator Pro", theme="soft") as demo:
    gr.Markdown("# 🎬 Subtitle Translator Pro")
    gr.Markdown("Dịch phụ đề (.srt) giữa nhiều ngôn ngữ khác nhau")
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload file phụ đề (.srt)", file_types=[".srt"])
            
            with gr.Row():
                source_lang = gr.Dropdown(
                    choices=list(LANGUAGE_MODELS.keys()),
                    value="Tiếng Anh",
                    label="Ngôn ngữ nguồn",
                    interactive=True
                )
                auto_detect = gr.Checkbox(label="Tự động phát hiện ngôn ngữ", value=True)
            
            target_lang = gr.Dropdown(
                choices=list(LANGUAGE_MODELS.keys()),
                value="Tiếng Việt",
                label="Ngôn ngữ đích"
            )
            
            with gr.Row():
                hours = gr.Number(label="Giờ", value=0, precision=0, minimum=0)
                minutes = gr.Number(label="Phút", value=0, precision=0, minimum=0, maximum=59)
                seconds = gr.Number(label="Giây", value=0, minimum=0, step=0.1)
            
            with gr.Row():
                add_time_btn = gr.Button("Chỉ thêm thời gian", variant="secondary")
                translate_btn = gr.Button("Dịch phụ đề", variant="primary")
        
        with gr.Column():
            file_output = gr.File(label="File phụ đề đã xử lý", interactive=False)
            status = gr.Textbox(label="Trạng thái")
    
    # Xử lý khi upload file
    def on_file_upload(file, auto_detect_flag):
        if file and auto_detect_flag:
            try:
                lang_code = detect_subtitle_language(file.name)
                detected_lang = LANGUAGE_CODES.get(lang_code, "Tiếng Anh")
                return gr.Dropdown(value=detected_lang)
            except:
                return gr.Dropdown(value="Tiếng Anh")
        return gr.Dropdown()
    
    file_input.upload(
        fn=on_file_upload,
        inputs=[file_input, auto_detect],
        outputs=source_lang
    )
    
    # Xử lý khi nhấn nút thêm thời gian
    add_time_btn.click(
        fn=add_time_to_subtitle,
        inputs=[file_input, hours, minutes, seconds],
        outputs=[file_output, status]
    )
    
    # Xử lý khi nhấn nút dịch phụ đề
    translate_btn.click(
        fn=translate_subtitle,
        inputs=[file_input, source_lang, target_lang, hours, minutes, seconds],
        outputs=[file_output, status]
    )
    
    gr.Markdown("### Thông tin")
    gr.Markdown("""
    - Hỗ trợ định dạng .srt
    - Tự động phát hiện ngôn ngữ nguồn
    - Dịch giữa 24 ngôn ngữ khác nhau
    - Hỗ trợ dịch qua tiếng Anh nếu không có model trực tiếp
    - Thêm thời gian vào tất cả phụ đề (hỗ trợ giây thập phân)
    - Có nút riêng để chỉ thêm thời gian trước khi dịch
    """)

if __name__ == "__main__":
    demo.launch()