mvsepless/__main__.py

import argparse
import gradio as gr
import os

if not __package__:
    from __init__ import mvsepless_app, Separator
else:
    from .__init__ import mvsepless_app, Separator

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MVSepless")
    subparsers = parser.add_subparsers(dest="command")
    app_parser = subparsers.add_parser("app", help="Приложение MVSepless")
    app_parser.add_argument(
        "--port", type=int, default=None, help="Порт для запуска сервера Gradio."
    )
    app_parser.add_argument(
        "--share",
        action="store_true",
        help="Создать публичную ссылку для приложения Gradio.",
    )
    cli_parser = subparsers.add_parser("cli", help="CLI MVSepless Lite")
    cli_parser.add_argument(
        "--input", type=str, required=True, help="Входной аудиофайл или каталог."
    )
    cli_parser.add_argument(
        "--output_dir", type=str, default=None, help="Каталог для выходных файлов."
    )
    cli_parser.add_argument(
        "--model_type",
        type=str,
        default="mel_band_roformer",
        help="Тип модели разделения.",
    )
    cli_parser.add_argument(
        "--model_name",
        type=str,
        default="Mel-Band-Roformer_Vocals_kimberley_jensen",
        help="Имя модели разделения.",
    )
    cli_parser.add_argument(
        "--ext_inst", action="store_true", help="Извлечь инструментал."
    )
    cli_parser.add_argument(
        "--output_format",
        type=str,
        default="mp3",
        choices=Separator.audio.output_formats,
        help="Формат выходного файла.",
    )
    cli_parser.add_argument(
        "--output_bitrate", type=str, default="320k", help="Битрейт выходного файла."
    )
    cli_parser.add_argument(
        "--template",
        type=str,
        default="NAME (STEM) MODEL",
        help="Шаблон именования выходных файлов.",
    )
    cli_parser.add_argument(
        "--selected_stems",
        type=str,
        nargs="*",
        default=None,
        help="Выбранные стемы для разделения.",
    )
    args = parser.parse_args()

    if args.command == "app":
        theme = gr.themes.Citrus(
            primary_hue="teal",
            secondary_hue="blue",
            neutral_hue="blue",
            spacing_size="sm",
            font=[
                gr.themes.GoogleFont("Montserrat"),
                "ui-sans-serif",
                "system-ui",
                "sans-serif",
            ],
        )
        mvsepless_lite_app = mvsepless_app(theme=theme)
        mvsepless_lite_app.launch(
            server_name="0.0.0.0",
            server_port=args.port,
            share=args.share,
            allowed_paths=["/"],
            debug=True,
        )
    elif args.command == "cli":
        input_data = args.input
        if os.path.isdir(input_data):
            list_valid_files = []
            for file in os.listdir(args.input):
                if os.path.isfile(os.path.join(args.input, file)):
                    if Separator.audio.check(os.path.join(args.input, file)):
                        list_valid_files.append(os.path.join(args.input, file))

            input_files = list_valid_files
        else:
            input_files = input_data

        results = Separator().separate(
            input=input_files,
            output_dir=args.output_dir,
            model_type=args.model_type,
            model_name=args.model_name,
            ext_inst=args.ext_inst,
            output_format=args.output_format,
            output_bitrate=args.output_bitrate,
            template=args.template,
            selected_stems=args.selected_stems,
        )
        print("Разделение завершено.")