Spaces:
Running on Zero
Running on Zero
| """Hugging Face ZeroGPU entry point for dots.tts. | |
| This Space runs an *unmodified*, pip-installed ``dots.tts`` package. All ZeroGPU | |
| specifics live in the local :mod:`dots_tts_zerogpu` patch layer: | |
| * the model source is never forked -- AOTI acceleration is applied by | |
| intercepting upstream's single compile seam; | |
| * generation buckets are warmed via upstream's own ``run_warmup``. | |
| Startup flow: | |
| 1. Apply the AOTI compile-seam patches (before any runtime is built). | |
| 2. Preload the runtime on the root process with warmup suppressed (root has only | |
| CUDA emulation, so no real compile can happen here). | |
| 3. Inside a ``@spaces.GPU`` call, run warmup under an AOTI capture session so | |
| every reachable step module is exported + AOTI-compiled, then hand the | |
| artifacts back and install them into the serving process. | |
| 4. Wrap ``service.generate`` with ``@spaces.GPU`` and launch Gradio. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import sys | |
| from pathlib import Path | |
| from typing import Any, Callable | |
| REPO_ROOT = Path(__file__).resolve().parent | |
| if str(REPO_ROOT) not in sys.path: | |
| sys.path.insert(0, str(REPO_ROOT)) | |
| class _SpacesFallback: | |
| """No-op ``spaces`` shim for local runs outside Hugging Face Spaces.""" | |
| def GPU(*decorator_args, **_decorator_kwargs): | |
| if decorator_args and callable(decorator_args[0]): | |
| return decorator_args[0] | |
| def decorate(fn: Callable[..., Any]) -> Callable[..., Any]: | |
| return fn | |
| return decorate | |
| try: | |
| import spaces # type: ignore | |
| except Exception: # pragma: no cover - only used outside Hugging Face Spaces. | |
| spaces = _SpacesFallback() # type: ignore | |
| def _env_bool(name: str, default: bool) -> bool: | |
| value = os.environ.get(name) | |
| if value is None: | |
| return default | |
| return value.strip().lower() in {"1", "true", "yes", "on"} | |
| def _env_int(name: str, default: int) -> int: | |
| value = os.environ.get(name) | |
| if value is None or not value.strip(): | |
| return default | |
| return int(value) | |
| def _configure_zero_gpu_environment() -> None: | |
| # optimize=True on CUDA would otherwise select the flex-attention backend, | |
| # which relies on module-level torch.compile that ZeroGPU cannot run. | |
| os.environ.setdefault("DOTS_TTS_DELAYED_DIT_BACKEND", "sdpa") | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| def _suppress_root_warmup() -> None: | |
| """Force ``DotsTtsRuntime.from_pretrained`` to skip the in-process warmup. | |
| The service loads the runtime lazily via ``from_pretrained`` and does not | |
| expose ``warmup_on_optimize``. On the root process there is no real GPU | |
| (only CUDA emulation), so any warmup here would run eagerly and slowly -- | |
| real compilation is done separately inside ``@spaces.GPU``. | |
| """ | |
| from dots_tts.runtime import DotsTtsRuntime | |
| original_from_pretrained = DotsTtsRuntime.from_pretrained.__func__ | |
| def _from_pretrained(cls, *args, **kwargs): | |
| kwargs.setdefault("warmup_on_optimize", False) | |
| return original_from_pretrained(cls, *args, **kwargs) | |
| DotsTtsRuntime.from_pretrained = classmethod(_from_pretrained) | |
| def _preload_runtime(app_service, app_config): | |
| runtime, resolved_model_name_or_path = app_service._get_runtime( # noqa: SLF001 | |
| app_config.default_model_name_or_path, | |
| ) | |
| return runtime, resolved_model_name_or_path | |
| def main() -> None: | |
| _configure_zero_gpu_environment() | |
| from loguru import logger | |
| from dots_tts_zerogpu import apply_zero_gpu_patches | |
| from dots_tts_zerogpu.aoti import AOTI_REGISTRY | |
| from dots_tts_zerogpu.warmup import compile_runtime_aoti | |
| # Patch the compile seam before any runtime (and thus any step module) exists. | |
| apply_zero_gpu_patches() | |
| _suppress_root_warmup() | |
| import gradio as gr | |
| from apps.gradio.app import PLAYGROUND_CSS, build_demo, build_playground_theme | |
| from apps.gradio.service import GradioAppService, build_gradio_app_config | |
| from dots_tts.utils.logging import configure_logging | |
| host = os.environ.get("DOTS_TTS_HOST", "0.0.0.0") | |
| port = _env_int("DOTS_TTS_PORT", 7860) | |
| model_name_or_path = os.environ.get( | |
| "DOTS_TTS_MODEL_NAME_OR_PATH", | |
| "dots-studio/dots.tts-soar", | |
| ) | |
| model_revision = os.environ.get("DOTS_TTS_MODEL_REVISION") or None | |
| precision = os.environ.get("DOTS_TTS_PRECISION", "bfloat16") | |
| execution_mode = os.environ.get("DOTS_TTS_EXECUTION_MODE", "generate") | |
| max_generate_length = _env_int("DOTS_TTS_MAX_GENERATE_LENGTH", 500) | |
| default_num_steps = _env_int("DOTS_TTS_DEFAULT_NUM_STEPS", 16) | |
| optimize = _env_bool("DOTS_TTS_OPTIMIZE", True) | |
| enable_aoti = _env_bool("DOTS_TTS_ENABLE_AOTI", True) | |
| startup_compile = _env_bool("DOTS_TTS_AOTI_COMPILE_ON_STARTUP", True) | |
| generation_duration = _env_int("DOTS_TTS_ZERO_GPU_DURATION", 60) | |
| compile_duration = _env_int("DOTS_TTS_ZERO_GPU_COMPILE_DURATION", 1500) | |
| output_dir = Path(os.environ.get("DOTS_TTS_OUTPUT_DIR", "/data/generated")) | |
| log_file = Path(os.environ.get("DOTS_TTS_LOG_FILE", "/tmp/dots_tts_gradio.log")) | |
| configure_logging(log_file=log_file) | |
| logger.info( | |
| "Space app starting: model={} execution_mode={} precision={} optimize={} " | |
| "enable_aoti={} startup_compile={} max_generate_length={}", | |
| model_name_or_path, | |
| execution_mode, | |
| precision, | |
| optimize, | |
| enable_aoti, | |
| startup_compile, | |
| max_generate_length, | |
| ) | |
| app_config = build_gradio_app_config( | |
| host=host, | |
| port=port, | |
| execution_mode=execution_mode, | |
| precision=precision, | |
| optimize=optimize, | |
| model_name_or_path=model_name_or_path, | |
| output_dir=output_dir, | |
| max_generate_length=max_generate_length, | |
| default_num_steps=default_num_steps, | |
| default_max_generate_length=max_generate_length, | |
| repo_root=REPO_ROOT, | |
| model_revision=model_revision, | |
| ) | |
| app_service = GradioAppService(app_config) | |
| runtime, resolved_model_name_or_path = _preload_runtime(app_service, app_config) | |
| if enable_aoti and startup_compile and optimize: | |
| def compile_aoti_cache(): | |
| child_runtime, _ = _preload_runtime(app_service, app_config) | |
| return compile_runtime_aoti(child_runtime) | |
| compiled_models = compile_aoti_cache() | |
| installed = AOTI_REGISTRY.install(compiled_models or {}) | |
| logger.info("AOTI startup compile completed: installed_target_count={}", installed) | |
| app_service.generate = spaces.GPU(duration=generation_duration)(app_service.generate) | |
| demo = build_demo(gr, app_config, app_service) | |
| logger.info( | |
| "Space app ready: host={} port={} resolved_model={} aoti_target_count={}", | |
| app_config.host, | |
| app_config.port, | |
| resolved_model_name_or_path, | |
| len(AOTI_REGISTRY), | |
| ) | |
| demo.launch( | |
| server_name=app_config.host, | |
| server_port=app_config.port, | |
| theme=build_playground_theme(gr), | |
| css=PLAYGROUND_CSS, | |
| ) | |
| if __name__ == "__main__": | |
| main() | |