lesiondiffusion-mlx 0.2.1__py3-none-any.whl

lesiondiffusion_mlx/__init__.py

@@ -0,0 +1,66 @@
+"""Native MLX implementation of LesionDiffusion, with lazy public imports."""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from lesiondiffusion_mlx._schedules import (
+        lesion_diffusion_timesteps as lesion_diffusion_timesteps,
+    )
+    from lesiondiffusion_mlx.mlx_model import (
+        LINET_CONFIG as LINET_CONFIG,
+    )
+    from lesiondiffusion_mlx.mlx_model import (
+        LMNET_CONFIG as LMNET_CONFIG,
+    )
+    from lesiondiffusion_mlx.mlx_model import (
+        BiomedTextEncoder as BiomedTextEncoder,
+    )
+    from lesiondiffusion_mlx.mlx_model import (
+        LdmUnet as LdmUnet,
+    )
+    from lesiondiffusion_mlx.mlx_model import (
+        SourceLayers as SourceLayers,
+    )
+    from lesiondiffusion_mlx.mlx_model import (
+        UnetConfig as UnetConfig,
+    )
+    from lesiondiffusion_mlx.mlx_model import (
+        VqAutoencoder as VqAutoencoder,
+    )
+    from lesiondiffusion_mlx.mlx_sampling import DdimSampler as DdimSampler
+    from lesiondiffusion_mlx.mlx_weights import write_conversion as write_conversion
+
+__version__ = "0.2.1"
+
+__all__ = [
+    "LINET_CONFIG",
+    "LMNET_CONFIG",
+    "BiomedTextEncoder",
+    "DdimSampler",
+    "LdmUnet",
+    "SourceLayers",
+    "UnetConfig",
+    "VqAutoencoder",
+    "lesion_diffusion_timesteps",
+    "write_conversion",
+]
+
+
+def __getattr__(name: str) -> object:
+    if name not in __all__:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module = (
+        "_schedules"
+        if name == "lesion_diffusion_timesteps"
+        else "mlx_weights"
+        if name == "write_conversion"
+        else "mlx_sampling"
+        if name == "DdimSampler"
+        else "mlx_model"
+    )
+    value = getattr(import_module(f"lesiondiffusion_mlx.{module}"), name)
+    globals()[name] = value
+    return value

lesiondiffusion_mlx/_arrays.py

@@ -0,0 +1,24 @@
+"""Host arrays shared by conversion, NIfTI I/O and reference tooling."""
+
+import numpy as np
+
+type HostArray = np.ndarray[tuple[int, ...], np.dtype[np.generic]]
+type FloatArray = np.ndarray[tuple[int, ...], np.dtype[np.float32]]
+type BoolArray = np.ndarray[tuple[int, ...], np.dtype[np.bool_]]
+type NumericScalar = (
+    np.bool_
+    | np.int8
+    | np.int16
+    | np.int32
+    | np.int64
+    | np.uint8
+    | np.uint16
+    | np.uint32
+    | np.uint64
+    | np.float16
+    | np.float32
+    | np.float64
+    | np.complex64
+    | np.complex128
+)
+type NumericArray = np.ndarray[tuple[int, ...], np.dtype[NumericScalar]]

lesiondiffusion_mlx/_nifti.py

@@ -0,0 +1,73 @@
+"""The NIfTI surface used here, verified against NiBabel's image/header APIs."""
+
+from os import PathLike
+from typing import Literal, Protocol, cast
+
+import nibabel
+import numpy as np
+from nibabel.filebasedimages import ImageFileError as ImageFileError
+
+from lesiondiffusion_mlx._arrays import HostArray
+
+type PathArg = str | PathLike[str]
+type SpatialAffine = np.ndarray[tuple[int, ...], np.dtype[np.float64]]
+
+
+class ArrayBuffer(Protocol):
+    def __array__(
+        self, dtype: np.dtype[np.generic] | None = None, copy: bool | None = None
+    ) -> HostArray: ...
+
+
+class NiftiHeader(Protocol):
+    def get_zooms(self) -> tuple[float, ...]: ...
+    def get_xyzt_units(self) -> tuple[str, str]: ...
+    def set_xyzt_units(self, xyz: str, t: str) -> None: ...
+    def set_data_dtype(self, dtype: type[np.generic]) -> None: ...
+
+
+class HeaderFactory(Protocol):
+    def from_header(self, header: NiftiHeader) -> NiftiHeader: ...
+
+
+class NiftiImage(Protocol):
+    shape: tuple[int, ...]
+    affine: SpatialAffine
+    dataobj: ArrayBuffer | HostArray
+    header: NiftiHeader
+
+    def get_qform(self, *, coded: Literal[True]) -> tuple[SpatialAffine | None, int]: ...
+    def get_sform(self, *, coded: Literal[True]) -> tuple[SpatialAffine | None, int]: ...
+    def set_qform(self, affine: HostArray, code: int) -> None: ...
+    def set_sform(self, affine: HostArray, code: int) -> None: ...
+
+
+class ImageFactory(Protocol):
+    def __call__(
+        self, dataobj: HostArray, affine: HostArray, header: NiftiHeader | None = None
+    ) -> NiftiImage: ...
+
+
+class OrientationFunctions(Protocol):
+    """The subset of ``nibabel.orientations`` used to resample to RAS 1 mm."""
+
+    def axcodes2ornt(self, axcodes: tuple[str, str, str]) -> HostArray: ...
+    def ornt_transform(self, start_ornt: HostArray, end_ornt: HostArray) -> HostArray: ...
+    def apply_orientation(self, data: HostArray, ornt: HostArray) -> HostArray: ...
+    def inv_ornt_aff(self, ornt: HostArray, shape: tuple[int, ...]) -> HostArray: ...
+
+
+class Nibabel(Protocol):
+    Nifti1Header: HeaderFactory
+    Nifti1Image: ImageFactory
+    orientations: OrientationFunctions
+
+    def load(self, filename: PathArg) -> NiftiImage: ...
+    def save(self, image: NiftiImage, filename: PathArg) -> None: ...
+    def aff2axcodes(self, affine: HostArray) -> tuple[str | None, ...]: ...
+    def io_orientation(self, affine: HostArray) -> HostArray: ...
+
+
+# This adapter is used only for .nii/.nii.gz inputs and NIfTI output; the wider
+# NiBabel loader supports unrelated image classes outside the runner contract.
+nib = cast(Nibabel, nibabel)

lesiondiffusion_mlx/_schedules.py

@@ -0,0 +1,38 @@
+"""Pinned upstream uniform DDIM timestep schedule, with explicit invalid-index errors."""
+
+from __future__ import annotations
+
+
+def _is_runtime_integer(value: object) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool)
+
+
+def lesion_diffusion_timesteps(steps: int, total: int = 1000) -> tuple[int, ...]:
+    """Return the pinned upstream uniform DDIM timestep indices, validated.
+
+    The indices are ``value + 1`` for ``value`` in ``range(0, total, total // steps)``,
+    in ascending order; the sampler walks them in reverse. The tuple has at least
+    ``steps`` entries and has more when ``steps`` does not divide ``total`` (7 entries
+    for ``steps=6``), so callers report ``len(...)`` as the actual step count. An
+    out-of-range index is rejected rather than clipped, to keep the schedule upstream's.
+
+    Args:
+        steps: Requested number of DDIM steps, a positive integer (not a bool).
+        total: Length of the training noise schedule (``alphas_cumprod``).
+
+    Raises:
+        ValueError: If ``steps`` or ``total`` is not an integer, ``steps < 1``,
+            ``total < steps``, or the last index would reach ``total`` (for example
+            ``steps=3`` or ``steps=1000`` with ``total=1000``).
+    """
+    if not _is_runtime_integer(steps) or steps < 1:
+        raise ValueError("LesionDiffusion DDIM steps must be a positive integer")
+    if not _is_runtime_integer(total) or total < steps:
+        raise ValueError("LesionDiffusion DDIM schedule must contain at least as many timesteps")
+    timesteps = tuple(value + 1 for value in range(0, total, total // steps))
+    if timesteps[-1] >= total:
+        raise ValueError(
+            f"LesionDiffusion DDIM steps {steps} produce out-of-bounds indices "
+            f"for a {total}-timestep schedule; choose another diffusion_steps value"
+        )
+    return timesteps