main.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
视频翻译管线 GUI 启动器
全局入口文件
"""
import atexit
import os
import signal
import sys
from pathlib import Path

# 添加项目根目录到 Python 路径
PROJECT_ROOT = Path(__file__).parent
sys.path.insert(0, str(PROJECT_ROOT))


def setup_model_cache_paths():
    """
    统一设置所有模型下载路径到项目 models/ 目录
    
    必须在导入任何 AI 库之前调用！
    
    设置的环境变量：
    - HF_HOME: HuggingFace 模型 (WhisperX, NLLB-200, RIFE, face_alignment)
    - TORCH_HOME: Torch Hub 模型 (Demucs)
    - COQUI_TOS_AGREED: Coqui TTS 同意条款（跳过交互提示）
    - TTS_HOME: Coqui TTS 模型目录 (XTTS v2)
    - PYTORCH_CUDA_ALLOC_CONF: CUDA 显存分配策略
    - PATH: 添加 FFmpeg 目录（供 whisper、demucs 等库使用）
    """
    models_dir = PROJECT_ROOT / "models"
    models_dir.mkdir(parents=True, exist_ok=True)
    
    # ========== FFmpeg 路径配置 ==========
    # 将配置文件中的 FFmpeg 目录添加到 PATH
    # 供第三方库（whisper.load_audio、demucs 等）直接调用
    from src.utils.ffmpeg_utils import get_ffmpeg_path
    ffmpeg_exe = get_ffmpeg_path()
    if ffmpeg_exe != "ffmpeg":
        ffmpeg_dir = Path(ffmpeg_exe).parent
        if ffmpeg_dir.exists():
            current_path = os.environ.get("PATH", "")
            if str(ffmpeg_dir) not in current_path:
                os.environ["PATH"] = str(ffmpeg_dir) + os.pathsep + current_path
                print(f"[FFmpeg] 已添加到 PATH: {ffmpeg_dir}")
    
    # HuggingFace 模型 (WhisperX, NLLB-200, RIFE 等)
    hf_home = models_dir / "huggingface"
    os.environ.setdefault("HF_HOME", str(hf_home))
    os.environ.setdefault("HUGGINGFACE_HUB_CACHE", str(hf_home / "hub"))
    
    # Torch Hub 模型 (Demucs, face_alignment 等)
    # ⚠️ 使用相对路径避免中文路径问题
    # PyTorch C++ JIT 后端无法正确处理包含中文字符的绝对路径
    # 使用相对路径 "models/torch" 可以正常工作
    os.environ.setdefault("TORCH_HOME", "models/torch")
    torch_home = models_dir / "torch"
    
    # Coqui TTS 模型 (XTTS v2)
    tts_home = models_dir / "tts"
    os.environ.setdefault("COQUI_TOS_AGREED", "1")  # 跳过条款确认
    os.environ.setdefault("TTS_HOME", str(tts_home))
    
    # CUDA 显存分配优化
    # expandable_segments: 允许 PyTorch 动态扩展显存段，减少碎片化
    # 这有助于在多模型切换场景（如 RIFE -> Wav2Lip）避免 OOM
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
    
    # 打印确认信息
    print(f"[模型缓存] HuggingFace: {hf_home}")
    print(f"[模型缓存] Torch Hub: {torch_home}")
    print(f"[模型缓存] Coqui TTS: {tts_home}")


# 🔑 在导入任何 AI 库之前设置缓存路径
setup_model_cache_paths()

from src.Gui import PipelineApplication


def cleanup_on_exit():
    """退出时清理资源"""
    print("\n[清理] 正在关闭所有资源...")
    
    # 1. 关闭 GPU 调度器
    try:
        from src.GpuScheduler import get_gpu_scheduler
        scheduler = get_gpu_scheduler()
        if scheduler:
            print("[清理] 关闭 GPU 调度器...")
            scheduler.shutdown(wait=False, timeout=2.0)
    except Exception as e:
        print(f"[清理] GPU 调度器关闭异常: {e}")
    
    # 2. 强制终止所有子进程
    try:
        import psutil
        current_process = psutil.Process()
        children = current_process.children(recursive=True)
        if children:
            print(f"[清理] 终止 {len(children)} 个子进程...")
            for child in children:
                try:
                    child.terminate()
                except psutil.NoSuchProcess:
                    pass
            
            # 等待最多2秒
            gone, alive = psutil.wait_procs(children, timeout=2)
            
            # 强制杀死还活着的进程
            for p in alive:
                try:
                    print(f"[清理] 强制终止进程: {p.pid}")
                    p.kill()
                except psutil.NoSuchProcess:
                    pass
    except ImportError:
        # psutil 不可用，使用 os 级别的清理
        pass
    except Exception as e:
        print(f"[清理] 子进程清理异常: {e}")
    
    print("[清理] 完成")


def main():
    """主函数"""
    # 注册退出清理函数
    atexit.register(cleanup_on_exit)
    
    # 注册信号处理（Ctrl+C 等）
    def signal_handler(signum, frame):
        print(f"\n[信号] 收到信号 {signum}，准备退出...")
        cleanup_on_exit()
        sys.exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    try:
        # 创建应用程序
        app = PipelineApplication()
        
        # 运行应用
        exit_code = app.run()
        
        # 正常退出前清理
        cleanup_on_exit()
        sys.exit(exit_code)
        
    except Exception as e:
        print(f"启动失败: {e}")
        import traceback
        traceback.print_exc()
        cleanup_on_exit()
        sys.exit(1)


if __name__ == "__main__":
    main()