pyproject.toml

[build-system]
requires = ["setuptools", "setuptools-scm"]
build-backend = "setuptools.build_meta"

[project]
name = "lightly_train"
description = "Train models with self-supervised learning in a single command"
authors = [
    { name = "Lightly Team" },
]
readme = "README.md"
license = { "text" = "AGPL-3.0" }
keywords = [
    "machine-learning",
    "computer-vision",
    "deep-learning",
    "self-supervised-learning",
    "contrastive-learning",
    "pytorch",
    "python",
    "pretrained-models",
    "embeddings",
]
dynamic = ["version"]
requires-python = ">=3.8, <3.14"  # Python 3.14 is not yet supported by Lightning
dependencies = [
    "albumentations>=1.3.1,<1.4.0; python_version=='3.8'",  # 1.3 is needed for SuperGradients. 1.4 has a lot of errors. >1.4 is not supported by Python 3.8.
    "albumentations>=1.3.1,!=1.4.*; python_version>='3.9'",   # 1.3 is needed for SuperGradients. 1.4 has a lot of errors.
    "eval-type-backport>=0.2",  # Required for Pydantic. Allows 3.10 types in old Python versions.
    "filelock>=3.0.0",     # Cross-platform file locking
    "fsspec>=2023.1.0",
    "lightly>=1.5.26",
    "matplotlib>=3.5",
    "omegaconf>=2.3",
    "psutil>=5.0",
    "pyarrow>=15.0.0",
    "nvidia-ml-py>=12", # For GPU monitoring
    "pytorch_lightning>=2.1",
    "tensorboard>=2.10.0",
    "torch>=2.1.0",    
    "torchmetrics[detection]>=0.8",    # Torchmetrics is needed for fine-tuning.
                            # 0.8 is required by SuperGradients
                            # 1.5.2 is the last version that supports Python 3.8
                            # 1.7 is recommended version
    "torchvision>=0.16.0",  # torchvision 0.15.0 was yanked
    "transformers>=4.46",   # Transformers is needed for fine-tuning.
                            # 4.46 is last version that supports Python 3.8.
                            # 4.51.3 is the default for EoMT
    "tqdm>=4.0.0",
    "pycocotools>=2.0.11; python_version>='3.9'",
    "rich>=13.0",
    "numpy>1.23.0; python_version>='3.12'",  # numpy<=1.23 fails to build on Python>3.12 (no distutils).
]

[dependency-groups]
dev = [
    "licenseheaders>=0.8.8",
    "mdformat>=0.7.21; python_version>='3.9'",  # Older versions change formatting behavior
    "mdformat-myst>=0.3.0; python_version>='3.10'", # MyST support for mdformat
    "mypy>=1.10.0",
    "pillow>=10.3.0,<12.0.0",  # helpers.create_image currently fails for newer pillow versions
    "pre-commit>=3.5.0",
    "types-pycocotools>=0.1.1; python_version>='3.9'",
    "pytest-examples>=0.0.13,<0.0.16", # 0.0.16 and 0.0.17 (latest to date) will cause the following failure when running pytest: "TypeError: 'ABCMeta' object is not subscriptable"
    "pytest-mock>=3.14.0",
    "pytest-split>=0.10.0; python_full_version>='3.8.1'",  # Test sharding in CI
    "pytest>=8.0.0",
    "ruff>=0.4.0",
    "types-PyYAML>=6.0.12",
    "types-tqdm>=4.0.0",
    "jinja2>=3.1.6",

    # Docs
    "furo>=2024.04.27",
    "myst_nb>=1.1.0; python_version>='3.9'",           # Jupyter notebook support (requires Python >=3.9)
    "myst_parser>=3.0",
    "sphinx-autoapi>=3.0",
    "sphinx-copybutton>=0.5.2",
    "sphinx-design>=0.6.1; python_version>='3.9'",
    "sphinx-inline-tabs>=2023.4.21",
    "sphinxcontrib-googleanalytics>=0.4",
    "sphinx>=7.1",              # Allow old sphinx version for Python 3.8 compatibility
]
dist = [
    "build>=1.2.2",
    "twine>=6.1.0",
    "wheel>=0.45.1",
]

# Pinned versions for Torch and TorchVision to avoid issues with the CUDA/driver version
# on the CI machine. These versions are compatible with CUDA 11.4 and Python 3.8.
# They are the latest versions available before 2024-08-28.
# Torch 2.5+ is no longer compatible with Python 3.8.
# Be careful when making changes on the CI machine as other repositories also depend on
# the installed CUDA/driver versions.
minimal-torch = [
    "torch==2.1.0; python_version<'3.9'",
    "torchvision==0.16.0; python_version<'3.9'",
]
pinned-torch-minimal = [
    "torch==2.4.0; python_version<'3.9'",
    "torchvision==0.19.0; python_version<'3.9'",
]
pinned-torch-maximal = [
    "torch==2.7.0; python_version>='3.9'",
    "torchvision==0.22.0; python_version>='3.9'",
]
pinned-rocm-torch = [
    "torch==2.4.1+rocm6.1; python_version<'3.10' and sys_platform=='linux' and platform_machine=='x86_64'",
    "torchvision==0.19.1+rocm6.1; python_version<'3.10' and sys_platform=='linux' and platform_machine=='x86_64'",
    "pytorch-triton-rocm==3.0.0; python_version<'3.10' and sys_platform=='linux' and platform_machine=='x86_64'",
    "torch==2.13.0+rocm7.2; python_version>='3.10' and sys_platform=='linux' and platform_machine=='x86_64'",
    "torchvision==0.28.0+rocm7.2; python_version>='3.10' and sys_platform=='linux' and platform_machine=='x86_64'",
    "triton-rocm; python_version>='3.10' and sys_platform=='linux' and platform_machine=='x86_64'",
]

[project.optional-dependencies]
bitsandbytes = [
    # Optional 8-bit optimizer state (optim_type="adamw8bit"). Imported lazily,
    # so the dependency is only required when that optimizer is selected.
    "bitsandbytes>=0.43.0",
]
dicom = [
    "pydicom>=2.1.0",
]
faster-coco-eval = [
    # Faster C++ COCOeval used for the mAP metric. Selected automatically for object
    # detection when installed, see _metrics/mean_average_precision.py.
    # Requires torchmetrics>=1.2, which added the `backend` argument.
    # 1.6.3 is the version torchmetrics itself tests against.
    # No wheels are built for macOS x86_64 and musllinux, therefore this is an
    # optional extra instead of a required dependency.
    "faster-coco-eval>=1.6.3; python_version>='3.9'",
]
mlflow = [
    "mlflow>=2.8.0", # 2.8.0 added stepped image logging (key/step args to log_image)
    "pandas>=2.2.3; python_version>='3.12'",  # pandas<2.2.3 fails to build on Python>3.12 (no pkg_resources).
]
notebook = [
    "notebook>=7.2.3",
]
# Note: rfdetr, or one of its dependencies, installs opencv-python and opencv-python-headless
# which can result in conflicts.
rfdetr = [
    "rfdetr>=1.4.1,<1.6.0; python_version>='3.10'", # TODO (Lionel, 03/26): Remove upper bound once unit test are fixed.
    "onnxruntime>=1.21.0; python_version>='3.10'", # onnxruntime is required for onnxsim, which is required for rfdetr
]
super-gradients = [
    "albumentations>=1.3.1,<1.4.0; python_version<'3.12'", # SuperGradients requires albumentations 1.3
    "super-gradients>=3.7; python_version<'3.12'",  # SuperGradients depends on numpy<=1.23 which requires distutils (removed in Python 3.12).
]
tensorboard = [] # Keep for backwards compatibility
timm = [
    "timm>=1.0.3",
]
ultralytics = [
    "ultralytics>=8.1.0",
]
wandb = [
    "wandb>=0.17.2", # contain fixes for numpy>=2.0 support
]
xformers = [
    "xformers>=0.0.18",
]
onnx = [
    "onnx>=1.15.0",
    "onnx>=1.18.0; python_version>='3.9'",  # onnx 1.18+ requires Python 3.9+
    "onnx>=1.20.0; python_version>='3.10'",  # no wheels built anymore for Python<3.10
    "onnxscript>=0.7.0; python_version>='3.9'",
    "netron>=9.0.0; python_version>='3.9'",
]
onnxruntime = [
    "onnxruntime>=1.21.0; python_version>='3.10'",
] # inference engines
onnxslim = [
    "onnxslim>=0.1.67",
]
[project.scripts]
lightly-train = "lightly_train._cli:_cli_entrypoint"

[tool.setuptools]
package-dir = {""= "src"}

[tool.setuptools.dynamic]
version = {attr = "lightly_train.__version__"}

[tool.mypy]
strict = true
disallow_untyped_defs = true
ignore_missing_imports = false
# Disabled because some type: ignore flags are only required on specific Python versions
# resulting in false positives.
warn_unused_ignores = false

[[tool.mypy.overrides]]
module = [
    "albumentations.*",
    "fsspec.*",
    "lightly.*",
    "matplotlib.*",
    "lightning_fabric.*",
    "lightning_fabric.loggers.logger",
    "pydicom.*",
    "super_gradients.*",
    "tensorboard.*",
    "timm.*",
    "torchmetrics.*",
    "torchvision.*",
    "transformers.*",
    "ultralytics.*",
    "onnxruntime.*",
]
ignore_missing_imports = true
 
 
[[tool.mypy.overrides]]
module = [
    "lightly_train._models.dinov2_vit.dinov2_vit_src.*",
    "lightly_train._models.dinov3.dinov3_src.*",
    "lightly_train._task_models.object_detection_components.*"
]
ignore_errors = true


[[tool.mypy.overrides]]
module = [
    "lightning_fabric.utilities.rank_zero",
    "rfdetr.*",
]
ignore_missing_imports = true
disable_error_code = ["attr-defined"]


# Required for TIMM which does no re-export the create_model function in timm/__init__.py
implicit_reexport = true

[tool.ruff.lint]
# Enable isort rules.
extend-select = ["I"]

[tool.ruff.lint.isort]
# Make sure that lightly_train is in the last import group.
# For details, see https://github.com/lightly-ai/lightly-train/pull/67
known-first-party = ["lightly_train"]

[tool.ruff.lint.pydocstyle]
# Use Google-style docstrings.
convention = "google"

[tool.ruff.lint.per-file-ignores]
# Ignore `E402` (import violations) in root `__init__.py`.
"src/lightly_train/__init__.py" = ["E402"]

[tool.uv]
override-dependencies = [
    # Override dependency from SuperGradients on sphinx<=4.1
    "sphinx>=7.1",
    # Floor for transitive onnx dependencies (e.g. from onnxslim) to prevent the
    # resolver from picking ancient versions like onnx==0.1 under --resolution=lowest-direct.
    # The onnx extra's Python-versioned lower bounds (>=1.18 for 3.9+, >=1.20 for 3.10+)
    # are intentionally left to the resolver; the conflict declaration below prevents
    # super-gradients' onnx==1.15.0 pin from pulling those down.
    "onnx>=1.15.0",
]
conflicts = [
    [ # conflicting due to incompatible versions of onnxruntime
        { extra = "super-gradients" },
        { extra = "onnxruntime" },
    ],
    [ # conflicting due to incompatible versions of onnxruntime
        { extra = "super-gradients" },
        { extra = "rfdetr" },
    ],
    [ # super-gradients requires onnx==1.15.0 which is incompatible with onnx>=1.18.0 for Python>=3.9
        { extra = "super-gradients" },
        { extra = "onnx" },
    ],
    [ # conflicting pinned versions of torch and torchvision for Python 3.8
        { group = "minimal-torch" },
        { group = "pinned-torch-minimal" },
        { group = "pinned-torch-maximal" },
        { group = "pinned-rocm-torch" },
    ],
]

[tool.uv.sources]
# The CI versions are pinned to specific indices to make sure that a version
# compatible with the CUDA driver of CI is installed.
torch = [
    { index = "pytorch-cu118", group = "pinned-torch-minimal", marker = "python_version<'3.9'" },
    { index = "pytorch-cu118", group = "pinned-torch-maximal", marker = "python_version>='3.9'" },
    { index = "pytorch-rocm-all", group = "pinned-rocm-torch", marker = "python_version<'3.10' and sys_platform=='linux' and platform_machine=='x86_64'" },
    { index = "pytorch-rocm-all", group = "pinned-rocm-torch", marker = "python_version>='3.10' and sys_platform=='linux' and platform_machine=='x86_64'" },
]
torchvision = [
    { index = "pytorch-cu118", group = "pinned-torch-minimal", marker = "python_version<'3.9'" },
    { index = "pytorch-cu118", group = "pinned-torch-maximal", marker = "python_version>='3.9'" },
    { index = "pytorch-rocm-all", group = "pinned-rocm-torch", marker = "python_version<'3.10' and sys_platform=='linux' and platform_machine=='x86_64'" },
    { index = "pytorch-rocm-all", group = "pinned-rocm-torch", marker = "python_version>='3.10' and sys_platform=='linux' and platform_machine=='x86_64'" },
]
pytorch-triton-rocm = [
    { index = "pytorch-rocm-all", group = "pinned-rocm-torch", marker = "python_version<'3.10' and sys_platform=='linux' and platform_machine=='x86_64'" },
]
triton-rocm = [
    { index = "pytorch-rocm-all", group = "pinned-rocm-torch", marker = "python_version>='3.10' and sys_platform=='linux' and platform_machine=='x86_64'" },
]

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

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

[tool.pytest.ini_options]
# Make pytest discover all files in the tests directory. This is usuful for showing
# detailed assertion errors outside of the test file, for example when using a helper
# function from `tests/helpers.py`.
python_files = "tests/*.py tests/**/*.py"
addopts = "--strict-markers"