RustyMark/dots.tts

1from __future__ import annotations2 3import os4import sys5from pathlib import Path6from typing import Any, Callable7 8 9REPO_ROOT = Path(__file__).resolve().parent10SRC_ROOT = REPO_ROOT / "src"11 12for import_root in (REPO_ROOT, SRC_ROOT):13    import_root_str = str(import_root)14    if import_root_str not in sys.path:15        sys.path.insert(0, import_root_str)16 17 18class _SpacesFallback:19    @staticmethod20    def GPU(*decorator_args, **_decorator_kwargs):21        if decorator_args and callable(decorator_args[0]):22            return decorator_args[0]23 24        def decorate(fn: Callable[..., Any]) -> Callable[..., Any]:25            return fn26 27        return decorate28 29 30try:31    import spaces  # type: ignore32except Exception:  # pragma: no cover - only used outside Hugging Face Spaces.33    spaces = _SpacesFallback()  # type: ignore34 35 36def _env_bool(name: str, default: bool) -> bool:37    value = os.environ.get(name)38    if value is None:39        return default40    return value.strip().lower() in {"1", "true", "yes", "on"}41 42 43def _env_int(name: str, default: int) -> int:44    value = os.environ.get(name)45    if value is None or not value.strip():46        return default47    return int(value)48 49 50def _configure_zero_gpu_environment() -> None:51    os.environ.setdefault("DOTS_TTS_COMPILE_BACKEND", "aoti")52    os.environ.setdefault("DOTS_TTS_SKIP_INIT_WARMUP", "1")53    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")54 55 56def _preload_runtime(app_service, app_config, compile_backend: str):57    runtime, resolved_model_name_or_path = app_service._get_runtime(  # noqa: SLF00158        app_config.default_model_name_or_path,59    )60    runtime.optimize = bool(app_config.optimize)61    runtime.model.set_optimize(bool(app_config.optimize))62    if hasattr(runtime.model, "set_compile_backend"):63        runtime.model.set_compile_backend(compile_backend)64    return runtime, resolved_model_name_or_path65 66 67def main() -> None:68    _configure_zero_gpu_environment()69 70    import gradio as gr71    from loguru import logger72 73    from apps.gradio.app import PLAYGROUND_CSS, build_demo, build_playground_theme74    from apps.gradio.service import GradioAppService, build_gradio_app_config75    from dots_tts.utils.logging import configure_logging76 77    host = os.environ.get("DOTS_TTS_HOST", "0.0.0.0")78    port = _env_int("DOTS_TTS_PORT", 7860)79    model_name_or_path = os.environ.get(80        "DOTS_TTS_MODEL_NAME_OR_PATH",81        "rednote-hilab/dots.tts",82    )83    model_revision = os.environ.get("DOTS_TTS_MODEL_REVISION") or None84    precision = os.environ.get("DOTS_TTS_PRECISION", "bfloat16")85    execution_mode = os.environ.get("DOTS_TTS_EXECUTION_MODE", "generate")86    max_generate_length = _env_int("DOTS_TTS_MAX_GENERATE_LENGTH", 500)87    default_num_steps = _env_int("DOTS_TTS_DEFAULT_NUM_STEPS", 16)88    compile_backend = os.environ.get("DOTS_TTS_COMPILE_BACKEND", "aoti").strip().lower()89    enable_aoti = _env_bool("DOTS_TTS_ENABLE_AOTI", True)90    startup_compile = _env_bool("DOTS_TTS_AOTI_COMPILE_ON_STARTUP", True)91    optimize = _env_bool("DOTS_TTS_OPTIMIZE", True)92    generation_duration = _env_int("DOTS_TTS_ZERO_GPU_DURATION", 60)93    compile_duration = _env_int("DOTS_TTS_ZERO_GPU_COMPILE_DURATION", 1500)94    output_dir = Path(os.environ.get("DOTS_TTS_OUTPUT_DIR", "/data/generated"))95    log_file = Path(os.environ.get("DOTS_TTS_LOG_FILE", "/tmp/dots_tts_gradio.log"))96 97    configure_logging(log_file=log_file)98    logger.info(99        "Space app starting: model={} execution_mode={} precision={} optimize={} "100        "compile_backend={} enable_aoti={} startup_compile={} max_generate_length={}",101        model_name_or_path,102        execution_mode,103        precision,104        optimize,105        compile_backend,106        enable_aoti,107        startup_compile,108        max_generate_length,109    )110 111    app_config = build_gradio_app_config(112        host=host,113        port=port,114        execution_mode=execution_mode,115        precision=precision,116        optimize=optimize,117        model_name_or_path=model_name_or_path,118        output_dir=output_dir,119        max_generate_length=max_generate_length,120        default_num_steps=default_num_steps,121        default_max_generate_length=max_generate_length,122        repo_root=REPO_ROOT,123        model_revision=model_revision,124    )125    app_service = GradioAppService(app_config)126    runtime, resolved_model_name_or_path = _preload_runtime(127        app_service,128        app_config,129        compile_backend if enable_aoti else "torch_compile",130    )131 132    if enable_aoti and startup_compile and optimize:133 134        @spaces.GPU(duration=compile_duration)135        def compile_aoti_cache():136            child_runtime, _ = _preload_runtime(137                app_service,138                app_config,139                compile_backend,140            )141            child_runtime.model.run_warmup(142                max_generate_length=app_config.max_generate_length,143                precision=app_config.precision,144                num_steps=app_config.default_num_steps,145                guidance_scale=app_config.default_guidance_scale,146            )147            return child_runtime.model.export_compiled_models()148 149        compiled_models = compile_aoti_cache()150        if compiled_models:151            runtime.model.import_compiled_models(compiled_models)152        logger.info(153            "AOTI startup compile completed: compiled_target_count={}",154            len(compiled_models or {}),155        )156 157    app_service.generate = spaces.GPU(duration=generation_duration)(app_service.generate)158 159    demo = build_demo(gr, app_config, app_service)160    logger.info(161        "Space app ready: host={} port={} resolved_model={} compiled_target_count={}",162        app_config.host,163        app_config.port,164        resolved_model_name_or_path,165        len(runtime.model.export_compiled_models())166        if hasattr(runtime.model, "export_compiled_models")167        else 0,168    )169    demo.launch(170        server_name=app_config.host,171        server_port=app_config.port,172        theme=build_playground_theme(gr),173        css=PLAYGROUND_CSS,174    )175 176 177if __name__ == "__main__":178    main()179