lesiondiffusion-mlx 0.2.1__py3-none-any.whl

lesiondiffusion_mlx/mlx_sampling.py

@@ -0,0 +1,132 @@
+"""DDIM sampler for the upstream mask and inpainting stages, executed on Metal."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, cast
+
+from lesiondiffusion_mlx._typing import ArrayConstructor
+
+if TYPE_CHECKING:
+    from lesiondiffusion_mlx._typing import Array, Mlx
+
+from lesiondiffusion_mlx._schedules import lesion_diffusion_timesteps
+
+
+class DdimSampler:
+    """Uniform DDIM sampler with upstream's schedule, covering both pipeline stages.
+
+    One class serves LMNet mask generation (no ``mask`` argument) and LINet latent
+    inpainting (with ``mask``), mirroring upstream's separate mask and inpainting
+    samplers. All tensors are channels-first ``(B, C, D, H, W)``. Timestep indices
+    come from ``lesion_diffusion_timesteps`` and are visited from the largest down.
+    ``eta=1`` (the default) matches upstream's stochastic sampling.
+    """
+
+    def __init__(self, alphas_cumprod: Array, mx: Mlx, steps: int, *, eta: float = 1.0):
+        """Precompute per-step coefficients.
+
+        Args:
+            alphas_cumprod: Cumulative product of the training noise schedule, shape
+                ``(T,)``, as stored in the converted checkpoint (float16 for LINet).
+            mx: The imported ``mlx.core`` module.
+            steps: Requested DDIM step count; see ``lesion_diffusion_timesteps``.
+            eta: DDIM stochasticity; 0 is deterministic, 1 is the upstream setting.
+
+        Raises:
+            ValueError: If ``steps`` does not give a valid schedule for ``T``.
+        """
+        self.mx = mx
+        self.timesteps = lesion_diffusion_timesteps(steps, alphas_cumprod.shape[0])
+        # Source LINet stores its schedule in fp16, then creates fp32 scalar
+        # broadcast tensors in p_sample_ddim. Preserve that rounded schedule.
+        alphas = alphas_cumprod[cast(ArrayConstructor, mx.array)(self.timesteps, dtype=mx.int32)]
+        previous = mx.concatenate([alphas_cumprod[:1], alphas[:-1]])
+        self.alphas = alphas.astype(mx.float32)
+        self.previous = previous.astype(mx.float32)
+        self.sqrt_one_minus = mx.sqrt(1 - alphas).astype(mx.float32)
+        # NumPy numerator / Torch denominator dispatches through Torch's
+        # reciprocal in the denominator dtype before promotion. This rounding
+        # is observable in LINet's fp16 schedule.
+        inverse_variance = mx.reciprocal(1 - alphas)
+        mx.eval(inverse_variance)
+        self.sigmas = eta * mx.sqrt(
+            (1 - self.previous)
+            * inverse_variance.astype(mx.float32)
+            * (1 - self.alphas / self.previous)
+        )
+        mx.eval(self.alphas, self.previous, self.sqrt_one_minus, self.sigmas)
+
+    def __call__(
+        self,
+        denoiser: Callable[[Array, Array, Array | None], Array],
+        condition: Array | None,
+        conditioning_image: Array,
+        *,
+        initial_noise: Array | None = None,
+        step_noise: tuple[Array, ...] | None = None,
+        mask: Array | None = None,
+    ) -> Array:
+        """Run the reverse process and return the final sample.
+
+        Without ``mask`` (LMNet), ``conditioning_image`` of shape ``(B, C, D, H, W)``
+        is concatenated to the noise on every step and the sample has shape
+        ``(B, 1, D, H, W)``. With ``mask`` (LINet), ``conditioning_image`` is the
+        clean latent ``(B, C, D, H, W)`` and ``mask`` is ``(B, 1, D, H, W)`` with 1
+        marking voxels to generate; voxels where ``mask`` is 0 are reset to
+        ``conditioning_image`` on every step and in the result, and the sample has
+        the shape of ``conditioning_image``.
+
+        Args:
+            denoiser: Network called as ``denoiser(input, timestep, condition)`` with
+                ``timestep`` an int32 array of shape ``(B,)``. Its output is the
+                noise prediction; in the masked case only the first ``C`` channels
+                are used and the network's extra mask channel is ignored.
+            condition: Cross-attention context ``(B, N, D_ctx)``, or ``None``.
+            conditioning_image: Conditioning volume or latent, see above.
+            initial_noise: Starting noise with the sample's shape; drawn from MLX's
+                RNG when omitted. Mainly for reference comparisons.
+            step_noise: One noise array per actual timestep, in the order the steps
+                run (largest timestep first); drawn from MLX's RNG when omitted.
+            mask: Region to generate, or ``None`` for unmasked sampling.
+
+        Returns:
+            The final sample. The lazy graph is evaluated after every step so it
+            does not grow across steps.
+
+        Raises:
+            ValueError: If ``initial_noise`` has the wrong shape or ``step_noise``
+                does not have one entry per timestep.
+        """
+        mx = self.mx
+        shape = (
+            (conditioning_image.shape[0], 1, *conditioning_image.shape[2:])
+            if mask is None
+            else conditioning_image.shape
+        )
+        x = mx.random.normal(shape) if initial_noise is None else initial_noise
+        if tuple(x.shape) != tuple(shape):
+            raise ValueError(f"DDIM initial noise shape {x.shape} differs from {shape}")
+        if step_noise is not None and len(step_noise) != len(self.timesteps):
+            raise ValueError("DDIM needs one noise tensor per actual source timestep")
+        for iteration, index in enumerate(reversed(range(len(self.timesteps)))):
+            timestep = mx.full((x.shape[0],), self.timesteps[index], dtype=mx.int32)
+            if mask is not None:
+                x = x * mask + (1 - mask) * conditioning_image
+                network_input = mx.concatenate([x, mask], axis=1)
+            else:
+                network_input = mx.concatenate([x, conditioning_image], axis=1)
+            eps = denoiser(network_input, timestep, condition)
+            if mask is not None:
+                # Source predicts a fifth mask channel but retains the supplied
+                # mask: p_sample_ddim's mask_pred is deliberately unused.
+                eps = eps[:, : conditioning_image.shape[1]]
+            alpha, previous, sigma = self.alphas[index], self.previous[index], self.sigmas[index]
+            predicted = (x - self.sqrt_one_minus[index] * eps) / mx.sqrt(alpha)
+            direction = mx.sqrt(1 - previous - sigma**2) * eps
+            noise = mx.random.normal(shape) if step_noise is None else step_noise[iteration]
+            x = mx.sqrt(previous) * predicted + direction + sigma * noise
+            mx.eval(x)
+        if mask is not None:
+            x = x * mask + (1 - mask) * conditioning_image
+        return x

lesiondiffusion_mlx/mlx_weights.py

@@ -0,0 +1,216 @@
+"""Checkpoint conversion without asset staging or runtime artifact loading.
+
+Torch deserialization is lazy and requires the optional ``conversion`` extra.
+Converted arrays retain source keys, NCDHW layouts, and float32 storage.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+import numpy as np
+from medmlx_core.conversion import load_torch_checkpoint
+
+from lesiondiffusion_mlx._arrays import FloatArray, HostArray
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+type StateTensor = HostArray | Tensor
+
+
+def denoiser_weights[T: StateTensor](state: Mapping[str, T]) -> dict[str, T]:
+    """Select the EMA parameters used by upstream's inference ``ema_scope``.
+
+    Each ``model.<name>`` entry (with any ``diffusion_model.`` prefix removed) is
+    replaced by ``model_ema.<name>``, where upstream's ``LitEma`` stores the key
+    with its dots removed. The returned keys are the source parameter names.
+
+    Raises:
+        ValueError: If an EMA entry is missing or differs in shape from its raw
+            parameter, or if the state has no denoiser parameters.
+    """
+    result: dict[str, T] = {}
+    for key, value in state.items():
+        if not key.startswith("model."):
+            continue
+        name = key.removeprefix("model.").removeprefix("diffusion_model.")
+        ema_name = "model_ema." + name.replace(".", "")
+        if ema_name not in state:
+            raise ValueError(f"Released LesionDiffusion checkpoint lacks EMA {ema_name}")
+        if value.shape != state[ema_name].shape:
+            raise ValueError(f"EMA shape differs for {name}")
+        result[name] = state[ema_name]
+    if not result:
+        raise ValueError("LesionDiffusion checkpoint has no denoiser parameters")
+    return result
+
+
+def convert_state_dict(
+    state: Mapping[str, StateTensor], *, stage: str
+) -> tuple[dict[str, FloatArray], dict[str, FloatArray]]:
+    """Return model and text arrays for ``lmnet``, ``linet``, or ``vq``.
+
+    Accept NumPy arrays or source Torch tensors. Release pinning, tokenizer
+    copying, NPZ writing, and conversion metadata belong to the caller.
+
+    Returns:
+        ``(arrays, text)``, both float32 NumPy arrays keyed by source name and in
+        source (NCDHW) layout. For ``lmnet`` and ``linet``, ``arrays`` holds the EMA
+        denoiser weights plus ``alphas_cumprod`` and ``text`` holds the BiomedCLIP
+        parameters under ``cond_stage_model.model.`` (prefix removed). For ``vq``,
+        ``arrays`` holds the encoder, decoder, quantizer and conv layers and ``text``
+        is empty.
+
+    Raises:
+        ValueError: For an unknown stage, missing EMA or text parameters, or
+            non-finite values.
+        KeyError: If a denoiser stage has no ``alphas_cumprod``.
+    """
+    if stage not in {"lmnet", "linet", "vq"}:
+        raise ValueError(f"Unknown LesionDiffusion stage: {stage}")
+    if stage == "vq":
+        selected = {
+            key: value
+            for key, value in state.items()
+            if key.startswith(
+                ("encoder.", "decoder.", "quant_conv.", "post_quant_conv.", "quantize.")
+            )
+        }
+    else:
+        selected = denoiser_weights(state)
+        selected["alphas_cumprod"] = state["alphas_cumprod"]
+    arrays = {
+        key: np.asarray(tensor, dtype=np.float32)
+        if isinstance(tensor, np.ndarray)
+        else cast(FloatArray, tensor.detach().float().numpy())
+        for key, tensor in selected.items()
+    }
+    filename = {"lmnet": "LMNet.ckpt", "linet": "LINet.ckpt", "vq": "VQ.ckpt"}[stage]
+    if any(not np.all(np.isfinite(value)) for value in arrays.values()):
+        raise ValueError(f"Non-finite weights in {filename}")
+    text: dict[str, FloatArray] = {}
+    if stage != "vq":
+        text_prefix = "cond_stage_model.model."
+        text = {
+            key.removeprefix(text_prefix): np.asarray(tensor, dtype=np.float32)
+            if isinstance(tensor, np.ndarray)
+            else cast(FloatArray, tensor.detach().float().numpy())
+            for key, tensor in state.items()
+            if key.startswith(text_prefix)
+        }
+        if not text or any(not np.all(np.isfinite(value)) for value in text.values()):
+            raise ValueError(f"Missing or non-finite text parameters in {filename}")
+    return arrays, text
+
+
+def convert_checkpoint(
+    path: Path, *, stage: str
+) -> tuple[dict[str, FloatArray], dict[str, FloatArray]]:
+    """Deserialize one explicit source checkpoint using the conversion extra.
+
+    Loads with ``weights_only=True`` and converts its ``state_dict``; see
+    ``convert_state_dict`` for the result and errors.
+    """
+    checkpoint = load_torch_checkpoint(path, weights_only=True)
+    return convert_state_dict(
+        cast(Mapping[str, StateTensor], checkpoint["state_dict"]), stage=stage
+    )
+
+
+def convert_checkpoints(checkpoints: Mapping[str, Path]) -> dict[str, dict[str, FloatArray]]:
+    """Convert explicit stage paths; return NPZ-ready mappings without staging files.
+
+    Keys of the result are ``lmnet``, ``linet`` and ``vq`` plus ``lmnet_text`` and
+    ``linet_text`` for the stages that carry a text encoder.
+    """
+    result: dict[str, dict[str, FloatArray]] = {}
+    for stage, path in checkpoints.items():
+        arrays, text = convert_checkpoint(path, stage=stage)
+        result[stage] = arrays
+        if stage != "vq":
+            result[stage + "_text"] = text
+    return result
+
+
+def write_conversion(
+    *,
+    checkpoints: Mapping[str, Path],
+    tokenizer_path: Path,
+    output_dir: Path,
+    overwrite: bool = False,
+) -> Path:
+    """Write a self-contained conversion for the MedMLX runner.
+
+    The architecture is pinned to ``reference.json``. The output directory gets one
+    float32 NPZ per stage (``lmnet``, ``linet``, ``vq``, ``lmnet_text``,
+    ``linet_text``), a copy of the tokenizer and ``conversion.json``, which records
+    the schema, upstream revision, layout and the EMA selection that ``run`` checks.
+
+    Args:
+        checkpoints: Upstream ``.ckpt`` paths keyed by exactly ``lmnet``, ``linet``
+            and ``vq``.
+        tokenizer_path: BiomedCLIP ``tokenizer.json`` copied alongside the weights.
+        output_dir: Destination directory, created if needed.
+        overwrite: Replace existing output files instead of refusing.
+
+    Returns:
+        The resolved output directory.
+
+    Raises:
+        InvalidInputError: If the checkpoint keys are wrong, an input is missing, an
+            output would overwrite an input or an existing file (without
+            ``overwrite``).
+        ValueError: If a checkpoint lacks EMA or text parameters or has non-finite
+            values; nothing is written in that case.
+    """
+    import json
+
+    from medmlx_core.errors import InvalidInputError
+
+    from lesiondiffusion_mlx.medmlx import (
+        CONVERSION_FILE,
+        CONVERSION_SCHEMA,
+        CONVERSION_VERSION,
+        CONVERTED_FILES,
+        SOURCE_REF,
+        TOKENIZER_FILE,
+    )
+
+    if set(checkpoints) != {"lmnet", "linet", "vq"}:
+        raise InvalidInputError("Conversion requires explicit lmnet, linet and vq checkpoints")
+    output = output_dir.resolve()
+    destinations = [output / name for name in (*CONVERTED_FILES, TOKENIZER_FILE, CONVERSION_FILE)]
+    sources = [path.resolve() for path in (*checkpoints.values(), tokenizer_path)]
+    if set(destinations) & set(sources):
+        raise InvalidInputError("Conversion output must not overwrite an input file")
+    if not overwrite and any(path.exists() for path in destinations):
+        raise InvalidInputError("Conversion output exists; pass overwrite=True to replace it")
+    for path in sources:
+        if not path.is_file():
+            raise InvalidInputError(f"Missing conversion input: {path}")
+    # Parse tokenizer before writing any output, including all model conversions.
+    from tokenizers import Tokenizer
+
+    Tokenizer.from_file(str(tokenizer_path))
+    stages = convert_checkpoints(checkpoints)
+    output.mkdir(parents=True, exist_ok=True)
+    for stage, arrays in stages.items():
+        # ``allow_pickle`` became a real savez parameter only in NumPy 2.2 (previously
+        # ``**kwds`` swallowed it as a spurious saved array named "allow_pickle"); the
+        # numpy floor in pyproject.toml is pinned to 2.2 for this reason.
+        np.savez(output / (stage + ".npz"), allow_pickle=False, **arrays)
+    (output / TOKENIZER_FILE).write_bytes(tokenizer_path.read_bytes())
+    metadata = {
+        "schema_version": CONVERSION_SCHEMA,
+        "converter_version": CONVERSION_VERSION,
+        "source_ref": SOURCE_REF,
+        "layout": "source_torch_ncdhw",
+        "stored_dtype": "float32",
+        "inference_dtype": {"lmnet": "float32", "linet": "float16", "vq": "float16"},
+        "ema": True,
+    }
+    (output / CONVERSION_FILE).write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
+    return output

lesiondiffusion_mlx/reference.json

@@ -0,0 +1,141 @@
+{
+  "schema_version": 1,
+  "repository": "HengruiTianSJTU/LesionDiffusion",
+  "revision": "880736e2b8a4d9becdb6594da7f8af6db19d04d1",
+  "text_repository": "mlfoundations/open_clip",
+  "text_revision": "3ff1faf10b60be27252be7f6c84ce7c8c5e14ec8",
+  "files": [
+    "pipeline/preprocess.py",
+    "open_clip/hf_model.py",
+    "open_clip/hf_configs.py",
+    "pipeline/inference/sample_legacy.py",
+    "pipeline/inference/models.py",
+    "pipeline/inference/utils.py",
+    "pipeline/ldm/util.py",
+    "pipeline/ldm/models/autoencoder.py",
+    "pipeline/ldm/modules/attention.py",
+    "pipeline/ldm/modules/ema.py",
+    "pipeline/ldm/data/data_process_func.py",
+    "pipeline/ldm/data/Torchio_contrast_dataloader.py",
+    "pipeline/ldm/modules/encoders/biomedclip.py",
+    "pipeline/ldm/modules/diffusionmodules/util.py",
+    "pipeline/ldm/modules/diffusionmodules/model.py",
+    "pipeline/ldm/modules/diffusionmodules/openaimodel.py",
+    "pipeline/ldm/modules/distributions/distributions.py",
+    "pipeline/ldm/models/diffusion/ddim.py",
+    "pipeline/ldm/models/diffusion/ddpm.py",
+    "pipeline/main.py",
+    "pipeline/test.py",
+    "pipeline/configs/clip/open_clip_config.json",
+    "pipeline/configs/diffusion/maskdiffusion_test.yaml",
+    "pipeline/configs/diffusion/lesiondiffusion_test.yaml",
+    "pipeline/configs/autoencoder/lesiondiffusion_vq.yaml",
+    "pipeline/organ_type.json"
+  ],
+  "organ_labels": {
+    "1": "spleen",
+    "2": "kidney right",
+    "3": "kidney left",
+    "4": "gallbladder",
+    "5": "liver",
+    "6": "stomach",
+    "7": "aorta",
+    "8": "inferior vena cava",
+    "9": "portal vein and splenic vein",
+    "10": "pancreas",
+    "11": "adrenal gland right",
+    "12": "adrenal gland left",
+    "13": "lung upper lobe left",
+    "14": "lung lower lobe left",
+    "15": "lung upper lobe right",
+    "16": "lung middle lobe right",
+    "17": "lung lower lobe right",
+    "18": "vertebrae L5",
+    "19": "vertebrae L4",
+    "20": "vertebrae L3",
+    "21": "vertebrae L2",
+    "22": "vertebrae L1",
+    "23": "vertebrae T12",
+    "24": "vertebrae T11",
+    "25": "vertebrae T10",
+    "26": "vertebrae T9",
+    "27": "vertebrae T8",
+    "28": "vertebrae T7",
+    "29": "vertebrae T6",
+    "30": "vertebrae T5",
+    "31": "vertebrae T4",
+    "32": "vertebrae T3",
+    "33": "vertebrae T2",
+    "34": "vertebrae T1",
+    "35": "vertebrae C7",
+    "36": "vertebrae C6",
+    "37": "vertebrae C5",
+    "38": "vertebrae C4",
+    "39": "vertebrae C3",
+    "40": "vertebrae C2",
+    "41": "vertebrae C1",
+    "42": "esophagus",
+    "43": "trachea",
+    "44": "heart myocardium",
+    "45": "heart atrium left",
+    "46": "heart ventricle left",
+    "47": "heart atrium right",
+    "48": "heart ventricle right",
+    "49": "pulmonary artery",
+    "50": "brain",
+    "51": "iliac artery left",
+    "52": "iliac artery right",
+    "53": "iliac vena left",
+    "54": "iliac vena right",
+    "55": "small bowel",
+    "56": "duodenum",
+    "57": "colon",
+    "58": "rib left 1",
+    "59": "rib left 2",
+    "60": "rib left 3",
+    "61": "rib left 4",
+    "62": "rib left 5",
+    "63": "rib left 6",
+    "64": "rib left 7",
+    "65": "rib left 8",
+    "66": "rib left 9",
+    "67": "rib left 10",
+    "68": "rib left 11",
+    "69": "rib left 12",
+    "70": "rib right 1",
+    "71": "rib right 2",
+    "72": "rib right 3",
+    "73": "rib right 4",
+    "74": "rib right 5",
+    "75": "rib right 6",
+    "76": "rib right 7",
+    "77": "rib right 8",
+    "78": "rib right 9",
+    "79": "rib right 10",
+    "80": "rib right 11",
+    "81": "rib right 12",
+    "82": "humerus left",
+    "83": "humerus right",
+    "84": "scapula left",
+    "85": "scapula right",
+    "86": "clavicula left",
+    "87": "clavicula right",
+    "88": "femur left",
+    "89": "femur right",
+    "90": "hip left",
+    "91": "hip right",
+    "92": "sacrum",
+    "93": "face",
+    "94": "gluteus maximus left",
+    "95": "gluteus maximus right",
+    "96": "gluteus medius left",
+    "97": "gluteus medius right",
+    "98": "gluteus minimus left",
+    "99": "gluteus minimus right",
+    "100": "autochthon left",
+    "101": "autochthon right",
+    "102": "iliopsoas left",
+    "103": "iliopsoas right",
+    "104": "urinary bladder"
+  }
+}

lesiondiffusion_mlx-0.2.1.dist-info/METADATA

@@ -0,0 +1,99 @@
+Metadata-Version: 2.5
+Name: lesiondiffusion-mlx
+Version: 0.2.1
+Summary: Native MLX implementation of LesionDiffusion
+Author: Joseph Sandoval, Alfredo Sandoval Jr.
+License-Expression: Apache-2.0 AND MIT AND BSD-3-Clause
+License-File: LICENSE
+License-File: THIRD_PARTY_NOTICES.md
+Classifier: Operating System :: MacOS :: MacOS X
+Requires-Python: <3.15,>=3.13
+Requires-Dist: medmlx-core==0.2.3
+Requires-Dist: nibabel<6,>=5.3
+Requires-Dist: numpy<3,>=2.2
+Requires-Dist: safetensors<1,>=0.6
+Requires-Dist: scipy<2,>=1.15
+Requires-Dist: tokenizers<1,>=0.21
+Provides-Extra: conversion
+Requires-Dist: torch<3,>=2.13; extra == 'conversion'
+Provides-Extra: dev
+Requires-Dist: einops<0.9,>=0.8.1; extra == 'dev'
+Requires-Dist: hypothesis<7,>=6.100; extra == 'dev'
+Requires-Dist: pillow<14,>=13; extra == 'dev'
+Requires-Dist: pyright==1.1.414; extra == 'dev'
+Requires-Dist: pytest<10,>=8.4; extra == 'dev'
+Requires-Dist: ruff<1,>=0.13; extra == 'dev'
+Requires-Dist: scipy-stubs<1.19,>=1.18.1.1; extra == 'dev'
+Requires-Dist: torch<2.14,>=2.13; extra == 'dev'
+Provides-Extra: reference
+Requires-Dist: einops==0.8.2; extra == 'reference'
+Requires-Dist: omegaconf==2.3.1; extra == 'reference'
+Requires-Dist: transformers==5.19.0; extra == 'reference'
+Description-Content-Type: text/markdown
+
+# lesiondiffusion-mlx
+
+Add synthetic lesions to CT with native MLX, with controls for organ, size, placement and phenotype.
+
+[Docs](docs/technical.md) · [MedMLX](https://github.com/MedMLX)
+
+![Liver CT before and after lesion generation; the contour marks the generated lesion.](https://raw.githubusercontent.com/MedMLX/lesiondiffusion-mlx/main/docs/showcase.png)
+
+*Liver CT before and after lesion generation; the contour marks the generated lesion.*
+
+| Input | Output |
+| :--- | :--- |
+| CT + aligned binary organ mask | Synthetic CT and lesion mask |
+
+**Apple Silicon · Metal · Python 3.13 / 3.14** · [uv](https://docs.astral.sh/uv/)
+
+## Quick start
+
+```bash
+uv venv --python 3.13
+source .venv/bin/activate
+repo="https://github.com/MedMLX/lesiondiffusion-mlx"
+uv pip install \
+  "lesiondiffusion-mlx[conversion] @ git+$repo" \
+  huggingface-hub
+```
+
+Get `LMNet.ckpt`, `LINet.ckpt` and `VQ.ckpt` from the [upstream weights folder](https://drive.google.com/drive/folders/1n6_eAWhsHBFBYtDxPtp_FXp_tNHhOCHi). Convert once, then generate:
+
+```python
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download
+from lesiondiffusion_mlx import write_conversion
+
+tokenizer = hf_hub_download("microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224", "tokenizer.json")
+write_conversion(
+    checkpoints={"lmnet": Path("LMNet.ckpt"), "linet": Path("LINet.ckpt"), "vq": Path("VQ.ckpt")},
+    tokenizer_path=Path(tokenizer),
+    output_dir=Path("lesiondiffusion-weights"),
+)
+```
+
+```python
+from lesiondiffusion_mlx.medmlx import run
+
+result = run(
+    image_path="ct_1mm_ras.nii.gz",
+    organ_mask_path="liver_mask.nii.gz",
+    weights_path="lesiondiffusion-weights",
+    output_dir="out",
+    organ="liver",
+    lesion_type="tumor",
+    bbox_size_mm=[30, 30, 20],
+    placement="inside_parenchyma",
+)
+print(result["outputs"]["synthetic_image"])
+```
+
+Inputs must support the documented crop sizes. Outputs use a RAS 1 mm grid; voxels outside the lesion can change. Synthetic output is not ground truth.
+
+## License
+
+[Apache-2.0](LICENSE). See [third-party notices](THIRD_PARTY_NOTICES.md). Model weights retain their publisher terms.
+
+For research and education. Outputs require human review.

lesiondiffusion_mlx-0.2.1.dist-info/RECORD

@@ -0,0 +1,20 @@
+lesiondiffusion_mlx/__init__.py,sha256=UeGgZA4aYNnDN9ehDIdnD_DGBrWBN90SS-2vUp9vjq8,1888
+lesiondiffusion_mlx/_arrays.py,sha256=80ywKgaNBlkOadJiEzvNPlFOd54yked_4nHFTJU7Rdo,617
+lesiondiffusion_mlx/_nifti.py,sha256=VctealMh5XsLDAThw3nsYNPOEUS7zYSrfJHOuunFUxY,2594
+lesiondiffusion_mlx/_schedules.py,sha256=kNOBm6o29lVWHjb44TCgdv0iXiNzZcQdShKcFq-42jQ,1866
+lesiondiffusion_mlx/_source_half.py,sha256=B6LA7xO-PPA9-O1imqa_cmkNi-BpcZFmDuxMCOVsOV8,20912
+lesiondiffusion_mlx/_spatial.py,sha256=sVtRgbWUg16YJhUAs8tWUqNmYr3sstzTENpT-JXzAsk,2015
+lesiondiffusion_mlx/_typing.py,sha256=rf2pV_VR4VeJK25UfVf1mcNzXQtNzJ2T7NwH9OrTJfI,3677
+lesiondiffusion_mlx/medmlx.py,sha256=C_nBE-6ZiwU4hb5OplRDNfPCwcb97ZS2IHK7Fer87cE,73568
+lesiondiffusion_mlx/mlx_model.py,sha256=k0pd4I31Oa_HbN0Yi5H90caoBxOpbTpBBgjIKVIIA1s,20282
+lesiondiffusion_mlx/mlx_pipeline.py,sha256=0RaG3wLv-ulpnBtWhNCORxZa1O53gT56X4QdpyLGxUY,5637
+lesiondiffusion_mlx/mlx_sampling.py,sha256=M3muk0AWKJ_ek7eK9uhrogmM28rcr-eRM662lIIuHy4,6650
+lesiondiffusion_mlx/mlx_weights.py,sha256=Rel5WjwreY1Wo8a3WKHRRXKYxRctN6L7MGjBYV0S4Rc,8946
+lesiondiffusion_mlx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lesiondiffusion_mlx/reference.json,sha256=tGAHnA3XwV4KrQQwZRyRzjKqyJ5Fx6jDvIEnZymRVd0,4210
+lesiondiffusion_mlx-0.2.1.dist-info/METADATA,sha256=Sutps01ZjmvxK-jLu0REurZGHoeHI0EKPmtBeV98xXw,3476
+lesiondiffusion_mlx-0.2.1.dist-info/WHEEL,sha256=2iZrJkDIkeDh0k9SikKMqFuRP-ua6mb4Ir1AsER4BEw,87
+lesiondiffusion_mlx-0.2.1.dist-info/entry_points.txt,sha256=9i2dc9GW_EHGf1XvevGK0W1rrfrsb8hKLq7HZ70io54,67
+lesiondiffusion_mlx-0.2.1.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+lesiondiffusion_mlx-0.2.1.dist-info/licenses/THIRD_PARTY_NOTICES.md,sha256=jsI1422mmHSWHxbHHU0p_JWF_sBfaga2XOKRlCEWMW0,5770
+lesiondiffusion_mlx-0.2.1.dist-info/RECORD,,

lesiondiffusion_mlx-0.2.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.33.0
+Root-Is-Purelib: true
+Tag: py3-none-any

lesiondiffusion_mlx-0.2.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[medmlx.models]
+lesiondiffusion = lesiondiffusion_mlx.medmlx:MODEL