[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "mbd_lm"
dynamic = ["version"]
description = "Multi Block Diffusion for Ultra Fast dLLM Inference"
readme = "README.md"
requires-python = ">=3.11, <3.13"
license = { file = "LICENSE" }
authors = [{ name = "Yijie Jin", email = "drewjin0827@gmail.com" }]

dependencies = [
  "blobfile>=3.0.0",
  "datasets>=2.16.0,<=2.21.0",
  "packaging>=23.0,<26.0",
  "torchdata>=0.8.0,<1.0",
  "tiktoken>=0.9.0",
  "diffusers>=0.30.0,<=0.31.0",
  "transformers==4.57.3",
  "psutil",
  "timm",
  "wandb",
  "setuptools",
  "einops>=0.8.1",
  "rich",
  "colorama",
  "build",
  "easydict>=1.13",
  "omegaconf>=2.3.0",
]

[project.optional-dependencies]
dev = [
  "pre-commit>=4.0.0,<5.0",
  "ruff>=0.7.0,<1.0",
  "pytest>=6.0.0,<8.0",
  "expecttest>=0.3.0,<0.4",
]
audio = ["av>=14.3.0,<15.0", "librosa>=0.11.0,<0.12", "soundfile>=0.13.1,<0.14"]
dit = ["diffusers>=0.30.0,<=0.31.0", "bitsandbytes>=0.46.0,<=0.47.0"]
npu_aarch64 = [
  "torch==2.7.1",
  "torch-npu==2.7.1",
  "torchvision==0.22.1",
  "torchaudio==2.7.1",
  "decorator>=5.2.1",
  "scipy>=1.16.2",
]
npu = [
  "torch==2.7.1+cpu",
  "torch-npu==2.7.1",
  "torchvision==0.22.1+cpu",
  "torchaudio==2.7.1+cpu",
  "decorator>=5.2.1",
  "scipy>=1.16.2",
]
gpu = [
  "lm-eval>=0.4.9.2",
  "accelerate>=1.12.0",
  "torch==2.8.0+cu128",
  "torchvision==0.23.0+cu128",
  "torchaudio==2.8.0+cu128",
  "nvidia-cusparselt-cu12",
  "nvidia-nccl-cu12",
  "liger-kernel",
  "flash-attn",
  "flash-attn-3",
  "hf_transfer",
  "veomni",
  "keep-gpu>=0.4.1",
  "nvitop>=1.6.1",
]
megatron = ["megatron-energon>=7.2.1"]
trl = ["trl<=0.9.6"]

[dependency-groups]
# Follow the best practice in https://docs.astral.sh/uv/concepts/projects/dependencies/#development-dependencies
# to manage dev dependencies (i.e., dependencies that are only used in development) like
# test, lint and doc tools.
dev = [{ include-group = "lint" }, { include-group = "test" }]
lint = ["pre-commit", "ruff"]
test = ["pytest", "expecttest"]

[tool.uv]
# This locks the uv version so that we have a consistent uv behavior across the board.
# Inconsistent uv versions might generate different uv lock files which creates chaos.
#
# NOTE 1: Update this at least once per month as uv releases new version every week.
# NOTE 2: When updating this line, make sure to update Dockerfile under docker/ to the same
#         version and release new docker images.
required-version = ">=0.8.14"
no-build-isolation-package = ["flash-attn", "flash-attn-3", "veomni"]

override-dependencies = [
  # Adding the (extra == 'gpu') mark to declare this override is only for gpu extra.
  # Without this mark, uv would download torch/torchaudio/torchvision unexpectedly from
  # PyPI without using the cu128 versions.
  # Ref: https://github.com/astral-sh/uv/issues/11153
  "torch==2.8.0+cu128; (extra == 'gpu')",
  "torchaudio==2.8.0+cu128; (extra == 'gpu')",
  "torchvision==0.23.0+cu128; (extra == 'gpu')",
  # Adding explicit override for NPU backends - aarch64 variant
  "torch==2.7.1; (extra == 'npu_aarch64')",
  "torchaudio==2.7.1; (extra == 'npu_aarch64')",
  "torchvision==0.22.1; (extra == 'npu_aarch64')",
  # Adding explicit override for NPU backends - x86 variant (default npu)
  "torch==2.7.1+cpu; (extra == 'npu')",
  "torchaudio==2.7.1+cpu; (extra == 'npu')",
  "torchvision==0.22.1+cpu; (extra == 'npu')",
  # Fixate to a specific transformers version since upgrading transformers is a large task.
  "transformers==4.57.3",
]

conflicts = [[{ extra = "gpu" }, { extra = "npu" }, { extra = "npu_aarch64" }]]

[tool.uv.sources]
# To support override-dependencies for multiple extras, one package can only have one index.
# Using pytorch index for all variants since it has both regular and +cpu versions
torch = [
  { index = "pytorch", marker = "(extra == 'npu' or extra == 'npu_aarch64')" },
  # Due to some unknown bugs, uv might resolve to a wheel from https://download.pytorch.org/whl/cu128_full which does
  # not include the proper dependencies on nvidia packages like nvidia-cublas and nvidia-nccl. This confuses uv to uninstall
  # those nvidia packages.
  #
  # By forcing it to the proper package here, we avoid the problem that uv could randomly resolve to other packages.
  #
  # NOTE: WHEN UPDATE TORCH VERSIONS, PLEASE UPDATE THIS TO THE CORRESPONDING WHEEL URL.
  { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp311-cp311-manylinux_2_28_x86_64.whl", extra = 'gpu' },
]
torchvision = { index = "pytorch" }
torchaudio = { index = "pytorch" }
# Download flash-attn wheel directly to avoid build issues.
flash-attn = { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp311-cp311-linux_x86_64.whl", marker = "extra == 'gpu' or extra == 'sglang'" }
flash-attn-3 = { url = "https://github.com/windreamer/flash-attention3-wheels/releases/download/2025.09.28/flash_attn_3-3.0.0b1%2B20250928.cu128torch280cxx11abitrue.5059fd-cp39-abi3-linux_x86_64.whl", marker = "extra == 'gpu'" }
veomni = { path = "VeOmni/dist/veomni-0.1.0-py3-none-any.whl", marker = "extra == 'gpu'" }

[[tool.uv.index]]
url = "https://mirrors.aliyun.com/pypi/simple"
default = true

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

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

[[tool.uv.dependency-metadata]]
name = "flash-attn"
requires-dist = ["torch"]

[[tool.uv.dependency-metadata]]
name = "flash-attn-3"
requires-dist = ["torch"]

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

[tool.setuptools.packages.find]
where = ["."]
include = ["mbd_lm"]

[tool.ruff]
target-version = "py311"
line-length = 119
indent-width = 4

[tool.ruff.lint]
ignore = ["C901", "E501", "E741", "W605", "C408"]
select = ["C", "E", "F", "I", "W"]

[tool.ruff.lint.per-file-ignores]
"__init__.py" = ["E402", "F401", "F403", "F811"]

[tool.ruff.lint.isort]
lines-after-imports = 2
known-first-party = ["mbd_lm"]
known-third-party = ["torch", "transformers", "wandb"]

[tool.ruff.format]
quote-style = "double"
indent-style = "space"
skip-magic-trailing-comma = false
line-ending = "auto"