pyproject.toml

[project]
name = "jflux"
version = "0.1.0"
description = "Inference codebase for Flux in Jax"
readme = "README.md"
requires-python = ">=3.10"
license = { file = "LICENSE.md" }
dependencies = [
    "fire>=0.7.0",
    "flax>=0.9.0",
]

[project.optional-dependencies]
gpu = [
  "jax[cuda]>=0.4.34",
]
tpu = [
  "libtpu-nightly @ https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20241017+nightly-py3-none-any.whl",
  "torch",
  "torchvision",
  "sentencepiece",
  "transformers",
  "einops"
]

[project.scripts]
jflux = "jflux.cli:app"

[tool.uv]
package = true
dev-dependencies = [
    "flux",
    "pytest>=8.3.3",
]

[tool.uv.sources]
jflux = { workspace = true }
flux = { git = "https://github.com/black-forest-labs/flux.git" }

[tool.ruff.lint]
select = ["E", "F", "I001", "W"]

[tool.ruff.lint.isort]
lines-after-imports = 2

[tool.ruff.lint.pydocstyle]
convention = "google"

[tool.mypy]
disable_error_code = "no-redef"

[tool.pytest.ini_options]
pythonpath = "."
testpaths = ["tests"]