[project]
name = "ace-step"
version = "1.5.0"
description = "ACE-Step 1.5"
readme = "README.md"
requires-python = ">=3.11,<3.13"
license = {text = "MIT"}
dependencies = [
    # PyTorch for Windows with CUDA 12.8
    "torch==2.7.1+cu128; sys_platform == 'win32'",
    "torchvision==0.22.1+cu128; sys_platform == 'win32'",
    "torchaudio==2.7.1+cu128; sys_platform == 'win32'",
    # PyTorch for Linux x86_64 with CUDA 12.8
    "torch==2.10.0+cu128; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "torchvision==0.25.0+cu128; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "torchaudio==2.10.0+cu128; sys_platform == 'linux' and platform_machine == 'x86_64'",
    # PyTorch for Linux aarch64 with CUDA 13.0 (e.g. NVIDIA DGX Spark)
    "torch==2.10.0+cu130; sys_platform == 'linux' and platform_machine == 'aarch64'",
    "torchvision==0.25.0+cu130; sys_platform == 'linux' and platform_machine == 'aarch64'",
    "torchaudio==2.10.0+cu130; sys_platform == 'linux' and platform_machine == 'aarch64'",
    # PyTorch for macOS (CPU / MPS)
    "torch>=2.9.1; sys_platform == 'darwin' and platform_machine == 'arm64'",
    "torchvision; sys_platform == 'darwin' and platform_machine == 'arm64'",
    "torchaudio>=2.9.1; sys_platform == 'darwin' and platform_machine == 'arm64'",
    # Common dependencies
    "transformers>=4.51.0,<4.58.0",
    "diffusers",
    "gradio==6.2.0",
    "matplotlib>=3.7.5",
    "scipy>=1.10.1",
    "soundfile>=0.13.1",
    "loguru>=0.7.3",
    "einops>=0.8.1",
    "accelerate>=1.12.0",
    "fastapi>=0.110.0",
    "diskcache",
    "uvicorn[standard]>=0.27.0",
    "numba>=0.63.1",
    "vector-quantize-pytorch>=1.27.15",
    "torchcodec>=0.9.1; platform_machine != 'aarch64'",
    "torchao>=0.14.1,<0.16.0; platform_machine != 'aarch64'",
    "torchao; platform_machine == 'aarch64'",
    "toml",
    # MLX dependencies (Apple Silicon native acceleration - macOS only)
    "mlx>=0.25.2; sys_platform == 'darwin' and platform_machine == 'arm64'",
    "mlx-lm>=0.20.0; sys_platform == 'darwin' and platform_machine == 'arm64'",
    # Training dependencies
    "peft>=0.18.0",
    "lycoris-lora",
    "lightning>=2.0.0",
    "tensorboard>=2.0.0",
    # Local third-party packages
    # nano-vllm source is configured in [tool.uv.sources]
    "nano-vllm; sys_platform != 'darwin' or platform_machine != 'arm64'",
    "modelscope",
    "tensorboard>=2.20.0",
    "typer-slim>=0.21.1",
]

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true

[tool.uv]
required-environments = [
    "sys_platform == 'win32' and platform_machine == 'AMD64'",
    "sys_platform == 'linux' and platform_machine == 'x86_64'",
    "sys_platform == 'linux' and platform_machine == 'aarch64'",
    "sys_platform == 'darwin' and platform_machine == 'arm64'",
]

[tool.uv.sources]
nano-vllm = { path = "acestep/third_parts/nano-vllm" }
torch = [
    { index = "pytorch-cu128", marker = "sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')" },
    { index = "pytorch-cu130", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
]
torchvision = [
    { index = "pytorch-cu128", marker = "sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')" },
    { index = "pytorch-cu130", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
]
torchaudio = [
    { index = "pytorch-cu128", marker = "sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')" },
    { index = "pytorch-cu130", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
]

[project.scripts]
acestep = "acestep.acestep_v15_pipeline:main"
acestep-api = "acestep.api_server:main"
acestep-download = "acestep.model_downloader:main"

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[dependency-groups]
dev = []

[tool.hatch.build.targets.wheel]
packages = ["acestep"]