setup.py

#!/usr/bin/env python3

# flake8: noqa: E501

import os
import platform
import re
import shutil
import subprocess
import sys
from pathlib import Path

from setuptools import find_packages, setup


def run_command(cmd):
    """Run a command and return True if successful, False otherwise"""
    try:
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        return result.returncode == 0
    except:
        return False


def detect_platform():
    """Detect platform using same logic as train.sh"""
    # Check for Darwin (macOS)
    if platform.system() == "Darwin":
        return "apple"

    # Check for NVIDIA GPU
    if run_command("nvidia-smi > /dev/null 2>&1"):
        return "cuda"

    # Check for ROCm
    if run_command("rocm-smi > /dev/null 2>&1"):
        return "rocm"

    # Check for ROCm environment variables (additional detection)
    if any(env in os.environ for env in ["ROCM_PATH", "HIP_PATH", "ROCM_HOME"]):
        return "rocm"

    # Check if ROCm tools exist
    if shutil.which("rocminfo") or shutil.which("rocm-smi"):
        return "rocm"

    # Default to CPU to avoid pulling incorrect GPU packages
    return "cpu"


def get_version():
    """Read version from simpletuner/__init__.py"""
    try:
        with open("simpletuner/__init__.py") as f:
            for line in f:
                if line.startswith("__version__"):
                    return line.split("=")[1].strip().strip('"').strip("'")
    except:
        pass
    return "3.0.0"


def _python_tag() -> str:
    """Return the CPython ABI tag for the current interpreter."""
    return f"cp{sys.version_info.major}{sys.version_info.minor}"


def _rocm_platform_tag():
    """Return the ROCm wheel platform tag, overridable via environment."""
    return os.environ.get("SIMPLETUNER_ROCM_PLATFORM_TAG", "linux_x86_64")


def _normalize_rocm_version(value: str) -> str:
    """Normalize ROCm version strings like 7.1.0 to 7.1 for wheel URLs."""
    match = re.search(r"(\d+)\.(\d+)", value)
    if match:
        return f"{match.group(1)}.{match.group(2)}"
    return value


def _strip_rocm_prefix(value: str) -> str:
    """Normalize ROCm tag overrides by removing common prefixes."""
    normalized = value.strip()
    if normalized.startswith("rocm-rel-"):
        return normalized[len("rocm-rel-") :]
    if normalized.startswith("rocm"):
        return normalized[len("rocm") :]
    return normalized


def _rocm_rel_version(rocm_version: str) -> str:
    """Return the ROCm release tag (major.minor.patch) used by repo.radeon.com."""
    override = os.environ.get("SIMPLETUNER_ROCM_REL") or os.environ.get("ROCM_REL")
    if override:
        return _strip_rocm_prefix(override)
    if rocm_version.startswith("7.1"):
        return "7.1.1"
    return rocm_version


def _rocm_build_tag(package: str, version: str, rocm_rel: str) -> str:
    """Return the per-package build tag suffix (e.g. lw.git351ff442)."""
    env_key = f"SIMPLETUNER_ROCM_{package.upper()}_BUILD_TAG"
    override = os.environ.get(env_key)
    if override:
        return override.lstrip(".")

    if rocm_rel == "7.1.1":
        defaults = {
            ("torch", "2.9.1"): "lw.git351ff442",
            ("torchvision", "0.24.0"): "gitb919bd0c",
            ("torchaudio", "2.9.0"): "gite3c6ee2b",
            ("triton", "3.5.1"): "gita272dfa8",
        }
        return defaults.get((package, version), "")

    return ""


def _rocm_wheel_tag(rocm_rel: str, build_tag: str = "") -> str:
    """Return the full rocm tag used in filenames."""
    override = os.environ.get("SIMPLETUNER_ROCM_WHEEL_TAG")
    if override:
        return _strip_rocm_prefix(override)
    if build_tag:
        return f"{rocm_rel}.{build_tag}"
    return rocm_rel


def _rocm_base_url(rocm_rel: str) -> str:
    """Return the base URL for ROCm wheels."""
    return os.environ.get(
        "SIMPLETUNER_ROCM_BASE_URL",
        f"https://repo.radeon.com/rocm/manylinux/rocm-rel-{rocm_rel}",
    )


def _detect_rocm_version() -> str:
    """Detect ROCm version from env or installed headers."""
    override = os.environ.get("SIMPLETUNER_ROCM_VERSION")
    if override:
        return _normalize_rocm_version(override)

    rocm_env = os.environ.get("ROCM_VERSION")
    if rocm_env:
        return _normalize_rocm_version(rocm_env)

    header_paths = [
        Path("/usr/include/rocm-core/rocm_version.h"),
        Path("/opt/rocm/include/rocm-core/rocm_version.h"),
        Path("/opt/rocm/include/rocm_version.h"),
    ]
    rocm_path = os.environ.get("ROCM_PATH")
    if rocm_path:
        header_paths.append(Path(rocm_path) / "include/rocm-core/rocm_version.h")
        header_paths.append(Path(rocm_path) / "include/rocm_version.h")

    for header_path in header_paths:
        try:
            content = header_path.read_text()
        except OSError:
            continue
        major = re.search(r"ROCM_VERSION_MAJOR\\s+(\\d+)", content)
        minor = re.search(r"ROCM_VERSION_MINOR\\s+(\\d+)", content)
        if major and minor:
            return f"{major.group(1)}.{minor.group(1)}"

    return "7.1"


def build_rocm_wheel_url(package: str, version: str, rocm_tag: str, base_url: str) -> str:
    """Build a direct wheel URL for ROCm packages."""
    py_tag = _python_tag()
    platform_tag = _rocm_platform_tag()
    filename = f"{package}-{version}+rocm{rocm_tag}-{py_tag}-{py_tag}-{platform_tag}.whl"
    return f"{package} @ {base_url}/{filename}"


def build_rocm_triton_wheel_url(triton_version: str, rocm_tag: str, base_url: str) -> str:
    """Build a direct wheel URL for Triton ROCm packages."""
    triton_base_url = os.environ.get("SIMPLETUNER_ROCM_TRITON_BASE_URL", base_url)
    return build_rocm_wheel_url("triton", triton_version, rocm_tag, triton_base_url)


def _resolve_ramtorch_dependency() -> str:
    """
    Prefer a local RamTorch checkout (default: ~/src/ramtorch) when present, otherwise fall back to the package name.
    """

    candidate_path = Path(os.environ.get("SIMPLETUNER_RAMTORCH_PATH", "~/src/ramtorch")).expanduser()
    try:
        if candidate_path.exists():
            return f"ramtorch @ {candidate_path.resolve().as_uri()}"
    except Exception:
        # Any failure falls back to the plain package spec.
        pass
    return "ramtorch"


def _cuda13_base_url() -> str:
    """Return the base URL for CUDA 13 PyTorch wheels."""
    return os.environ.get(
        "SIMPLETUNER_CUDA13_BASE_URL",
        "https://download.pytorch.org/whl/cu130",
    )


def _cuda_nightly_base_url() -> str:
    """Return the base URL for CUDA 12 nightly PyTorch wheels."""
    return os.environ.get(
        "SIMPLETUNER_CUDA_NIGHTLY_BASE_URL",
        "https://download.pytorch.org/whl/nightly/cu126",
    )


def _cuda13_nightly_base_url() -> str:
    """Return the base URL for CUDA 13 nightly PyTorch wheels."""
    return os.environ.get(
        "SIMPLETUNER_CUDA13_NIGHTLY_BASE_URL",
        "https://download.pytorch.org/whl/nightly/cu130",
    )


def build_cuda13_wheel_url(package: str, version: str) -> str:
    """Build a direct wheel URL for CUDA 13 PyTorch packages."""
    py_tag = _python_tag()
    base_url = _cuda13_base_url()
    platform_tag = os.environ.get("SIMPLETUNER_CUDA13_PLATFORM_TAG", "manylinux_2_28_x86_64")
    filename = f"{package}-{version}%2Bcu130-{py_tag}-{py_tag}-{platform_tag}.whl"
    return f"{package} @ {base_url}/{filename}"


def build_cuda_nightly_wheel_url(package: str, version: str) -> str:
    """Build a direct wheel URL for CUDA 12 nightly PyTorch packages."""
    py_tag = _python_tag()
    base_url = _cuda_nightly_base_url()
    platform_tag = os.environ.get("SIMPLETUNER_CUDA_NIGHTLY_PLATFORM_TAG", "manylinux_2_28_x86_64")
    filename = f"{package}-{version}%2Bcu126-{py_tag}-{py_tag}-{platform_tag}.whl"
    return f"{package} @ {base_url}/{filename}"


def build_cuda13_nightly_wheel_url(package: str, version: str) -> str:
    """Build a direct wheel URL for CUDA 13 nightly PyTorch packages."""
    py_tag = _python_tag()
    base_url = _cuda13_nightly_base_url()
    platform_tag = os.environ.get("SIMPLETUNER_CUDA13_NIGHTLY_PLATFORM_TAG", "manylinux_2_28_x86_64")
    filename = f"{package}-{version}%2Bcu130-{py_tag}-{py_tag}-{platform_tag}.whl"
    return f"{package} @ {base_url}/{filename}"


def build_triton_wheel_url(version: str, base_url: str) -> str:
    """Build a direct wheel URL for triton from PyTorch wheel indices."""
    py_tag = _python_tag()
    platform_tag = os.environ.get("SIMPLETUNER_TRITON_PLATFORM_TAG", "manylinux_2_27_x86_64.manylinux_2_28_x86_64")
    filename = f"triton-{version}-{py_tag}-{py_tag}-{platform_tag}.whl"
    return f"triton @ {base_url}/{filename}"


def get_cuda13_dependencies():
    """Get CUDA 13 specific dependencies with direct wheel URLs."""
    ramtorch_dep = _resolve_ramtorch_dependency()
    torch_version = os.environ.get("SIMPLETUNER_CUDA13_TORCH_VERSION", "2.10.0")
    torchvision_version = os.environ.get("SIMPLETUNER_CUDA13_TORCHVISION_VERSION", "0.25.0")
    torchaudio_version = os.environ.get("SIMPLETUNER_CUDA13_TORCHAUDIO_VERSION", "2.10.0")

    return [
        build_cuda13_wheel_url("torch", torch_version),
        build_cuda13_wheel_url("torchvision", torchvision_version),
        build_cuda13_wheel_url("torchaudio", torchaudio_version),
        "triton>=3.3.0",
        "deepspeed>=0.17.2",
        "torchao>=0.14.1",
        "bitsandbytes>=0.45.0",
        "nvidia-cudnn-cu13",
        "nvidia-nccl-cu13",
        "nvidia-ml-py>=12.555",
        "lm-eval>=0.4.4",
        ramtorch_dep,
    ]


def get_cuda_nightly_dependencies():
    """Get CUDA 12 nightly dependencies (PyTorch 2.11.0.dev) with direct wheel URLs."""
    ramtorch_dep = _resolve_ramtorch_dependency()
    torch_version = os.environ.get("SIMPLETUNER_CUDA_NIGHTLY_TORCH_VERSION", "2.11.0.dev20260201")
    torchvision_version = os.environ.get("SIMPLETUNER_CUDA_NIGHTLY_TORCHVISION_VERSION", "0.25.0.dev20260201")
    torchaudio_version = os.environ.get("SIMPLETUNER_CUDA_NIGHTLY_TORCHAUDIO_VERSION", "2.11.0.dev20260201")
    triton_version = os.environ.get("SIMPLETUNER_CUDA_NIGHTLY_TRITON_VERSION", "3.6.0+git9844da95")

    return [
        build_cuda_nightly_wheel_url("torch", torch_version),
        build_cuda_nightly_wheel_url("torchvision", torchvision_version),
        build_cuda_nightly_wheel_url("torchaudio", torchaudio_version),
        build_triton_wheel_url(triton_version, "https://download.pytorch.org/whl/nightly"),
        "bitsandbytes>=0.45.0",
        "deepspeed>=0.17.2",
        "torchao>=0.14.1",
        "nvidia-cudnn-cu12",
        "nvidia-nccl-cu12",
        "nvidia-ml-py>=12.555",
        "lm-eval>=0.4.4",
        ramtorch_dep,
    ]


def get_cuda13_nightly_dependencies():
    """Get CUDA 13 nightly dependencies (PyTorch 2.11.0.dev) with direct wheel URLs."""
    ramtorch_dep = _resolve_ramtorch_dependency()
    torch_version = os.environ.get("SIMPLETUNER_CUDA13_NIGHTLY_TORCH_VERSION", "2.11.0.dev20260201")
    torchvision_version = os.environ.get("SIMPLETUNER_CUDA13_NIGHTLY_TORCHVISION_VERSION", "0.25.0.dev20260201")
    torchaudio_version = os.environ.get("SIMPLETUNER_CUDA13_NIGHTLY_TORCHAUDIO_VERSION", "2.11.0.dev20260131")
    triton_version = os.environ.get("SIMPLETUNER_CUDA13_NIGHTLY_TRITON_VERSION", "3.6.0+git9844da95")

    return [
        build_cuda13_nightly_wheel_url("torch", torch_version),
        build_cuda13_nightly_wheel_url("torchvision", torchvision_version),
        build_cuda13_nightly_wheel_url("torchaudio", torchaudio_version),
        build_triton_wheel_url(triton_version, "https://download.pytorch.org/whl/nightly"),
        "deepspeed>=0.17.2",
        "torchao>=0.14.1",
        "bitsandbytes>=0.45.0",
        "nvidia-cudnn-cu13",
        "nvidia-nccl-cu13",
        "nvidia-ml-py>=12.555",
        "lm-eval>=0.4.4",
        ramtorch_dep,
    ]


def get_cuda_dependencies():
    ramtorch_dep = _resolve_ramtorch_dependency()
    return [
        "torch>=2.10.0",
        "torchvision>=0.25.0",
        "torchaudio>=2.10.0",
        "triton>=3.3.0",
        "bitsandbytes>=0.45.0",
        "deepspeed>=0.17.2",
        "torchao>=0.14.1",
        "nvidia-cudnn-cu12",
        "nvidia-nccl-cu12",
        "nvidia-ml-py>=12.555",
        "lm-eval>=0.4.4",
        ramtorch_dep,
    ]


def get_rocm_dependencies():
    ramtorch_dep = _resolve_ramtorch_dependency()
    rocm_version = _detect_rocm_version()
    rocm_rel = _rocm_rel_version(rocm_version)
    rocm_base_url = _rocm_base_url(rocm_rel)
    torch_version = os.environ.get("SIMPLETUNER_ROCM_TORCH_VERSION", "2.9.1")
    torchvision_version = os.environ.get("SIMPLETUNER_ROCM_TORCHVISION_VERSION", "0.24.0")
    torchaudio_version = os.environ.get("SIMPLETUNER_ROCM_TORCHAUDIO_VERSION", "2.9.0")
    triton_version = os.environ.get("SIMPLETUNER_ROCM_TRITON_VERSION", "3.5.1")
    torch_tag = _rocm_wheel_tag(rocm_rel, _rocm_build_tag("torch", torch_version, rocm_rel))
    vision_tag = _rocm_wheel_tag(rocm_rel, _rocm_build_tag("torchvision", torchvision_version, rocm_rel))
    audio_tag = _rocm_wheel_tag(rocm_rel, _rocm_build_tag("torchaudio", torchaudio_version, rocm_rel))
    triton_tag = _rocm_wheel_tag(rocm_rel, _rocm_build_tag("triton", triton_version, rocm_rel))

    try:
        return [
            build_rocm_wheel_url("torch", torch_version, torch_tag, rocm_base_url),
            build_rocm_wheel_url("torchvision", torchvision_version, vision_tag, rocm_base_url),
            build_rocm_wheel_url("torchaudio", torchaudio_version, audio_tag, rocm_base_url),
            build_rocm_triton_wheel_url(triton_version, triton_tag, rocm_base_url),
            "torchao>=0.14.1",
            ramtorch_dep,
        ]
    except Exception as exc:
        print(f"Warning: falling back to CPU PyTorch packages because ROCm wheel configuration failed: {exc}")
        return [
            "torch>=2.10.0",
            "torchvision>=0.25.0",
            "torchaudio>=2.10.0",
            "torchao>=0.14.1",
            ramtorch_dep,
        ]


def get_apple_dependencies():
    return [
        "torch>=2.10.0",
        "torchvision>=0.25.0",
        "torchaudio>=2.10.0",
        "torchao>=0.14.1",
    ]


def get_cpu_dependencies():
    return [
        "torch>=2.10.0",
        "torchvision>=0.25.0",
        "torchaudio>=2.10.0",
        "torchao>=0.14.1",
    ]


PLATFORM_DEPENDENCIES = {
    "cuda": get_cuda_dependencies(),
    "rocm": get_rocm_dependencies(),
    "apple": get_apple_dependencies(),
    "cpu": get_cpu_dependencies(),
}


def get_platform_dependencies():
    """Get dependencies based on detected platform"""
    detected_platform = detect_platform()

    # Allow override via environment variable
    platform_override = os.environ.get("SIMPLETUNER_PLATFORM", detected_platform)

    print(f"Detected platform: {detected_platform}")
    if platform_override != detected_platform:
        print(f"Platform overridden to: {platform_override}")

    platform_to_use = platform_override

    # Base PyTorch dependencies
    deps = PLATFORM_DEPENDENCIES.get(platform_to_use, PLATFORM_DEPENDENCIES["cpu"])
    print(f"Installing {platform_to_use.upper()} dependencies...")
    return deps


def _collect_package_files(*directories: str):
    """Collect package data files relative to the simpletuner package."""
    collected = []
    package_root = Path("simpletuner")
    for directory in directories:
        root = Path(directory)
        if not root.exists():
            continue
        for path in root.rglob("*"):
            if path.is_file():
                try:
                    relative = path.relative_to(package_root)
                except ValueError:
                    # Skip files outside package root
                    continue
                collected.append(str(relative))
    return collected


# Base dependencies (minimal, works on all platforms)
base_deps = [
    "diffusers>=0.36.0",
    "transformers>=4.55.0",
    "hf_transfer>=0.1.0",
    "datasets>=3.0.1",
    "wandb>=0.21.0",
    "requests>=2.32.4",
    "pillow>=11.3.0",
    "trainingsample>=0.2.10",
    "accelerate>=1.5.2",
    "safetensors>=0.5.3",
    "compel>=2.1.1",
    "clip-interrogator>=0.6.0",
    "open-clip-torch>=2.26.1",
    "iterutils>=0.1.6",
    "scipy>=1.11.1",
    "boto3>=1.35.83",
    "pandas>=2.2.3",
    "botocore>=1.35.83",
    "skrample>=0.5.0",
    "urllib3<1.27",
    "torchsde>=0.2.6",
    "torchmetrics>=1.1.1",
    "colorama>=0.4.6",
    "numpy>=2.2.0",
    "num2words>=0.5.13",
    "peft>=0.17.0",
    "tensorboard>=2.18.0",
    "py3langid>=0.2.2",
    "pypinyin>=0.50.0",
    "sentencepiece>=0.2.0",
    "spacy>=3.7.4",
    "hangul-romanize>=0.1.0",
    "optimum-quanto>=0.2.7",
    "lycoris-lora>=3.4.0",
    "torch-optimi>=0.2.1",
    "librosa>=0.10.2",
    "loguru>=0.7.2",
    "toml>=0.10.2",
    "fastapi[standard]>=0.115.0",
    "sse-starlette>=1.6.5",
    "atomicwrites>=1.4.1",
    "beautifulsoup4>=4.12.3",
    "prodigy-plus-schedule-free>=1.9.2",
    "tokenizers>=0.21.0",
    "huggingface-hub>=0.34.3",
    "imageio-ffmpeg>=0.6.0",
    "imageio[pyav]>=2.37.0",
    "hf-xet>=1.1.5",
    "peft-singlora>=0.2.0",
    "vector-quantize-pytorch>=1.27.15",
    "cryptography>=41.0.0",
    "torchcodec>=0.8.1",
    "sdnq>=0.1.2",
    "aiosqlite>=0.19.0",
    "httpx>=0.28.0",
    "psutil>=5.9.0",
]

# Optional extras
extras_require = {
    "jxl": ["pillow-jxl-plugin>=1.3.1"],
    "dev": [
        "selenium>=4.0.0",
        "coverage>=7.0.0",
        "black>=23.0.0",
        "isort>=5.12.0",
        "flake8>=6.0.0",
        "mypy>=1.0.0",
        "pre-commit>=3.0.0",
    ],
    "test": ["selenium>=4.0.0", "coverage>=7.0.0"],
    "docs": [
        "zensical>=0.0.19",
    ],
    # Platform-specific extras - user must choose one
    "cuda": list(PLATFORM_DEPENDENCIES["cuda"]),
    "cuda13": get_cuda13_dependencies(),
    "cuda-nightly": get_cuda_nightly_dependencies(),
    "cuda13-nightly": get_cuda13_nightly_dependencies(),
    "rocm": list(PLATFORM_DEPENDENCIES["rocm"]),
    "apple": list(PLATFORM_DEPENDENCIES["apple"]),
    "cpu": list(PLATFORM_DEPENDENCIES["cpu"]),
    # State backend extras for multi-node deployments
    "state-postgresql": ["asyncpg>=0.29.0"],
    "state-mysql": ["aiomysql>=0.2.0"],
    "state-redis": ["redis>=5.0.0"],
    "state-all": ["asyncpg>=0.29.0", "aiomysql>=0.2.0", "redis>=5.0.0"],
    # All non-platform extras combined
    "all": [
        "pillow-jxl-plugin>=1.3.1",
        "selenium>=4.0.0",
        "coverage>=7.0.0",
        "black>=23.0.0",
        "isort>=5.12.0",
        "flake8>=6.0.0",
    ],
}

# Read long description
try:
    with open("README.md", "r", encoding="utf-8") as fh:
        long_description = fh.read()
except:
    long_description = "Stable Diffusion 2.x and XL tuner."

setup(
    name="simpletuner",
    version=get_version(),
    description="Stable Diffusion 2.x and XL tuner.",
    long_description=long_description,
    long_description_content_type="text/markdown",
    author="bghira",
    # license handled by pyproject.toml
    packages=find_packages(),
    include_package_data=True,
    package_data={
        "simpletuner": _collect_package_files(
            "simpletuner/templates",
            "simpletuner/static",
            "simpletuner/config",
            "simpletuner/documentation",
        ),
    },
    python_requires=">=3.12,<3.14",
    install_requires=base_deps,
    extras_require=extras_require,
    entry_points={
        "console_scripts": [
            "simpletuner=simpletuner.cli:main",
            "simpletuner-train=simpletuner.train:main",
            "simpletuner-configure=simpletuner.configure:main",
            "simpletuner-inference=simpletuner.inference:main",
            "simpletuner-server=simpletuner.service_worker:main",
        ],
    },
    classifiers=[
        "Development Status :: 4 - Beta",
        "Intended Audience :: Developers",
        "License :: OSI Approved :: GNU Affero General Public License v3",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.12",
        "Programming Language :: Python :: 3.13",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Topic :: Multimedia :: Graphics",
    ],
    keywords="stable-diffusion machine-learning deep-learning pytorch cuda rocm",
    url="https://github.com/bghira/SimpleTuner",
    project_urls={
        "Bug Reports": "https://github.com/bghira/SimpleTuner/issues",
        "Source": "https://github.com/bghira/SimpleTuner",
        "Documentation": "https://github.com/bghira/SimpleTuner/blob/main/README.md",
    },
)

if __name__ == "__main__":
    print("SimpleTuner Setup")
    print("================")
    print(f"Detected platform: {detect_platform()}")
    print(f"Python version: {sys.version}")
    print(f"Platform: {platform.platform()}")
    print("\nInstall with a platform extra:")
    print("  pip install .[cuda]          # CUDA 12 (PyTorch 2.10.0 release)")
    print("  pip install .[cuda13]        # CUDA 13 (PyTorch 2.10.0 release)")
    print("  pip install .[cuda-nightly]  # CUDA 12 (PyTorch 2.11.0 nightly)")
    print("  pip install .[cuda13-nightly]# CUDA 13 (PyTorch 2.11.0 nightly)")
    print("  pip install .[rocm]          # ROCm")
    print("  pip install .[apple]         # macOS")
    print("  pip install .[cpu]           # CPU only")