# Copyright 2026 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2

EAPI=8

DISTUTILS_USE_PEP517=setuptools
DISTUTILS_SINGLE_IMPL=1
PYTHON_COMPAT=( python3_{12..14} )

inherit distutils-r1

DESCRIPTION="Fast kernel library for Diffusion inference with multiple compute backends"
HOMEPAGE="https://github.com/Comfy-Org/comfy-kitchen"
SRC_URI="https://github.com/Comfy-Org/${PN}/archive/refs/tags/v${PV}.tar.gz -> ${P}.gh.tar.gz"

S="${WORKDIR}/${P}"

LICENSE="Apache-2.0"
SLOT="0"
KEYWORDS="~amd64"

# No CUDA USE flag — we always build CPU-only for Gentoo.
# The triton and eager backends work for both CUDA and ROCm at runtime.
# The compiled CUDA extension requires CUDA 12.8+ and nanobind and is not
# needed for standard ComfyUI operation.

RDEPEND="
	|| (
		sci-ml/pytorch[${PYTHON_SINGLE_USEDEP}]
		sci-ml/pytorch-rocm[${PYTHON_SINGLE_USEDEP}]
		sci-ml/pytorch-cuda[${PYTHON_SINGLE_USEDEP}]
	)
"

src_prepare() {
	default
	# Allow env-var control of CUDA build (original only checks --no-cuda argv)
	sed -i 's/^BUILD_NO_CUDA = False/BUILD_NO_CUDA = os.environ.get("COMFY_KITCHEN_BUILD_NO_CUDA", "0") == "1"/' \
		setup.py || die
}

python_compile() {
	COMFY_KITCHEN_BUILD_NO_CUDA=1 distutils-r1_python_compile
}