Update v0.2.1 inference API and bundle metadata

Retain the exact converted tensors and tokenizer, preserve private archive and upstream component terms, and remove the superseded loader patch.

README.md

@@ -17,134 +17,72 @@ pipeline_tag: image-to-image
 
 # LesionDiffusion for MLX
 
-Native MLX LesionDiffusion synthesis on Apple Silicon using LMNet, LINet, and VQ components. The recorded paper configuration inserts a liver tumor into a prepared CT crop; current conditioned inference is documented separately below.
+Converted LMNet, LINet, text-encoder and VQ weights for native MLX lesion synthesis on Apple Silicon. This entry remains **private**; original checkpoint and component terms remain applicable. Intended for research and education, not clinical use.
 
-**Repository contents:** converted FP32 safetensors weights, runtime JSON metadata where required, and retained publisher notices. This repository is private; original checkpoint and component terms remain applicable.
+Use the public [lesiondiffusion-mlx v0.2.1 package](https://github.com/MedMLX/lesiondiffusion-mlx/releases/tag/v0.2.1), source revision `41f7400dd6f34b1b5844732ad2360a1301cdb9fa`, on an Apple Silicon Mac with Metal and Python 3.13 or 3.14. Inference does not require PyTorch. The package reads this consolidated safetensors bundle directly; no source patch is needed.
 
-Requires an Apple Silicon Mac with Metal and Python 3.12/3.13. Model inference uses MLX; PyTorch is used for one-time checkpoint conversion and optional upstream comparisons. Intended for research and education; clinical diagnostic performance has not been established.
+## Install and download
 
-Native MLX package: [lesiondiffusion-mlx](https://github.com/MedMLX/lesiondiffusion-mlx).
-Evaluated revision: `9b7a15eef808a38ab56e82325c22f209f69b9b78`.
-
-This card describes the one configuration evaluated in the MedMLX paper. It is research software; the single-case comparison does not estimate clinical accuracy.
-
-## Source and checkpoint
-
-Original publisher: HengruiTianSJTU.
-Original implementation: [https://github.com/HengruiTianSJTU/LesionDiffusion](https://github.com/HengruiTianSJTU/LesionDiffusion).
-Source revision or version: `880736e2b8a4d9becdb6594da7f8af6db19d04d1`.
-Original checkpoint: Component locations are in the native run record and the upstream package documentation.
-
-Conversion: lesiondiffusion_mlx.write_conversion (README recipe), using the upstream component checkpoints and an explicitly supplied BiomedCLIP tokenizer.
-Conversion changes tensor layout or storage as specified by the native package; it does not retrain the model.
-
-## Download the safetensors bundle
-
-Install [uv](https://docs.astral.sh/uv/getting-started/installation/) and Git first. Your Hugging Face account needs access to this private MedMLX entry, and your Git credentials need access to the companion private MedMLX source repository. In a new working directory, create and activate a Python environment, install the Hub CLI, and download the bundle:
+Install [uv](https://docs.astral.sh/uv/getting-started/installation/) and Git, then create a Python environment and install the pinned public package. Install the Hub CLI as a standalone tool; your Hugging Face account needs access to this private model entry.
 
 ```bash
-uv venv --python 3.12
+uv venv --python 3.13
 source .venv/bin/activate
-uv pip install huggingface_hub
+uv pip install "lesiondiffusion-mlx @ git+https://github.com/MedMLX/lesiondiffusion-mlx.git@41f7400dd6f34b1b5844732ad2360a1301cdb9fa"
+uv tool install huggingface-hub
 hf auth login
-hf download MedMLX/lesiondiffusion-mlx --local-dir lesiondiffusion
+hf download MedMLX/lesiondiffusion-mlx --local-dir lesiondiffusion-weights
 ```
 
-The current source checkout needs the included safetensors loader patch. The patch changes serialization support only; the existing local NPZ input contract is retained. Install the pinned source with that patch:
+The download contains `model.safetensors`, `conversion.json`, `tokenizer.json`, and retained notices. Weights are stored in FP32 with `<stage>.` prefixes: `lmnet`, `lmnet_text`, `linet`, `linet_text`, and `vq`. The package applies its declared inference precision profile when loading each stage.
 
-```bash
-git clone https://github.com/MedMLX/lesiondiffusion-mlx.git lesiondiffusion-mlx-source
-git -C lesiondiffusion-mlx-source checkout c8cac6e616895618254069bcf10ddf292409bc62
-git -C lesiondiffusion-mlx-source apply ../lesiondiffusion/safetensors-support.patch
-uv pip install ./lesiondiffusion-mlx-source
-```
-
-The recorded usage example below uses this bundle. The original conversion recipe is retained for reproducibility.
+## Public API
 
-## Reproduce the upstream conversion (optional)
-
-Install the native package with conversion support:
-
-```bash
-uv pip install "lesiondiffusion-mlx[conversion] @ git+https://github.com/MedMLX/lesiondiffusion-mlx.git"
-```
-
-Obtain the LMNet, LINet, and VQ checkpoints through the [official upstream project](https://github.com/HengruiTianSJTU/LesionDiffusion), together with its BiomedCLIP `tokenizer.json`. Component terms remain separate. Convert the explicitly selected local assets:
-
-```python
-from pathlib import Path
-from lesiondiffusion_mlx import write_conversion
-
-write_conversion(
-    checkpoints={
-        "lmnet": Path("LMNet.ckpt"),
-        "linet": Path("LINet.ckpt"),
-        "vq": Path("VQ.ckpt"),
-    },
-    tokenizer_path=Path("tokenizer.json"),
-    output_dir=Path("lesiondiffusion"),
-)
-```
-
-The [package README](https://github.com/MedMLX/lesiondiffusion-mlx#input-and-output) defines the CT and organ-mask preparation contract. The example below explicitly selects the paper's `source` inference mode. The package currently defaults to `conditioned`; its source-parity result must not be applied to that mode.
-
-## Recorded example
-
-AMOS22 (validation) / amos_0008; AMOS22 (validation) / amos_0008 label 6 (liver)
+Provide a CT NIfTI and an aligned binary organ mask (0 background, 1 organ). For a liver example:
 
 ```python
 from lesiondiffusion_mlx.medmlx import run
 
-# Replace file and directory names with your locally obtained inputs and weights.
-kwargs = {
-    "image_path": "amos_0008_ct_1mm_ras_160x160x128.nii.gz",
-    "organ_mask_path": "amos_0008_liver_1mm_ras_160x160x128.nii.gz",
-    "output_dir": "out",
-    "weights_path": "lesiondiffusion",
-    "organ": "liver",
-    "lesion_type": "tumor",
-    "bbox_size_mm": [
-        30,
-        30,
-        20
-    ],
-    "placement": "inside_parenchyma",
-    "inference_mode": "source",
-    "diffusion_steps": 200,
-    "overwrite": False,
-    "seed": 23
-}
-result = run(**kwargs)
-print(result["outputs"])
+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",
+    inference_mode="conditioned",
+    diffusion_steps=200,
+    seed=23,
+)
+print(result["outputs"]["synthetic_image"])
 ```
 
-Native result: **success**. Original-code comparison: **pass**, `cosine_similarity` = `0.9999697715207947`.
-
-Full controlled 200-step paired reference completed October 12; image cosine 0.99996977 and lesion Dice 1.0. Density and outside-lesion reconstruction limitations remain.
+The default `conditioned` mode sends phenotype text to LINet and uses native FP16 convolution. It changes generated outputs relative to `inference_mode="source"`, which follows the pinned upstream inference behavior, including disabled LINet image-text conditioning and source-ordered convolution arithmetic. Both modes default to 200 steps and seed 23. Source-mode agreement does not establish conditioned-mode equivalence or generated-image quality.
 
-The complete [native run](https://github.com/MedMLX/Med-MLX/blob/main/results/runs/lesiondiffusion.json) and [original-code comparison](https://github.com/MedMLX/Med-MLX/blob/main/results/reference/lesiondiffusion.json) retain commands, geometry, precision, and output-specific checks. Use the [replay harness](https://github.com/MedMLX/Med-MLX/blob/main/benchmarks/README.md) to substitute local paths. The package README documents installation and conversion.
+The [v0.2.1 input/output contract](https://github.com/MedMLX/lesiondiffusion-mlx/blob/v0.2.1/README.md#input-and-output) defines supported geometry, organ labels, placement and optional phenotype controls. Axis-aligned inputs are standardized to an even-dimensioned RAS 1 mm grid; oblique or sheared transforms are refused. Inputs must support the documented 152 × 152 × 64 and 128³ crops. The CT and mask must resolve to the same standardized grid.
 
-The current package defaults to `inference_mode="conditioned"`, which enables LINet text context and native FP16 convolution. The example above explicitly selects `source` to retain the paper's evaluated mode. The [October 12 follow-up](https://github.com/MedMLX/Med-MLX/blob/main/results/lesiondiffusion-conditioned-20261012.json) records a 452.6 s 200-step run with −20.6 HU lesion-minus-nearby-liver contrast and a 128.4 s 50-step draw with −19.0 HU contrast on the same case. This additional visibility check does not extend the paper's source-parity result to conditioned mode or establish clinical validity.
+A successful run returns output paths under `result["outputs"]`: `synthetic_image`, `target_mask` and `organ_mask`. The files are `synthetic_ct.nii.gz`, `lesion_mask.nii.gz` and `organ_mask.nii.gz` on the standardized grid. Existing outputs require explicit `overwrite=True` to replace them. In the default full reconstruction, voxels outside the lesion can also change, and CT intensities are clipped to [-1000, 1000] HU. See the package contract for `output_region="box"`. Synthetic output is not ground truth.
 
-## Distribution terms
+## Qualification scope
 
-Upstream code: **MIT**. Weights: **Unresolved; retrieve directly from the original publisher**.
-Distribution: **private converted-weight archive**. Public redistribution review remains separate; the original component terms are unchanged.
+The v0.2.1 release acceptance on October 24, 2026 used the installed wheel with public `medmlx-core==0.2.3`, MLX 0.32.4 and no PyTorch in the inference environment. A complete default conditioned 200-step, seed-23 public-API run finished in 372.0 seconds on one fixed AMOS22 `amos_0008` CT/liver crop (label 6; 160 × 160 × 128, RAS, 1 mm). The generated CT was finite, output masks had the required label types, and output geometry and coded forms matched the prepared input. This is an execution and output-contract check, not an assessment of appearance, lesion fidelity or clinical effectiveness.
 
-Full controlled 200-step paired reference completed October 12; image cosine 0.99996977 and lesion Dice 1.0. Density and outside-lesion reconstruction limitations remain. The complete component terms have not been cleared for public redistribution.
+All 139 current owner tests passed with zero skips, including the independent pinned-source numerical/contract tests. Their reduced-width and source-mode scope is described in the [technical notes](https://github.com/MedMLX/lesiondiffusion-mlx/blob/v0.2.1/docs/technical.md). The full conditioned run above is not a paired full-checkpoint source comparison.
 
-Required notices:
+Before this card update, every tensor in this Hub bundle was compared on CPU against the fresh canonical conversion used for that release acceptance: all 1,141 tensor keys, shapes, dtypes and values across five stages matched exactly, as did the tokenizer bytes. Both `source` and `conditioned` asset admission passed through the installed v0.2.1 public package. The Hub serialization was not given a separate full inference rerun; its tensor-equivalence check does not extend the model qualification.
 
-- Retain the applicable license, copyright and attribution notices; identify the upstream checkpoint, authors, and MLX conversion changes.
+## Source and conversion
 
-Review sources:
+Original publisher: HengruiTianSJTU. Source: [LesionDiffusion](https://github.com/HengruiTianSJTU/LesionDiffusion), pinned revision `880736e2b8a4d9becdb6594da7f8af6db19d04d1`.
 
-- https://github.com/HengruiTianSJTU/LesionDiffusion
+The bundle was converted from the released `LMNet.ckpt`, `LINet.ckpt` and `VQ.ckpt` components, with an explicitly supplied BiomedCLIP tokenizer. Conversion uses EMA weights and changes tensor layout/storage without retraining. `conversion.json` records the source revision, layout and precision profile. To repeat conversion from assets obtained from the original publisher, use the [v0.2.1 conversion recipe](https://github.com/MedMLX/lesiondiffusion-mlx/blob/v0.2.1/README.md#quick-start) and its optional `conversion` extra. Do not use this card as a grant to obtain or redistribute upstream assets.
 
-## Packaged weights
+## Distribution terms and notices
 
-Exactly one safetensors file contains the converted FP32 weights. JSON files are required runtime metadata or tokenizer assets. This upload does not rerun or extend the recorded model qualification.
+Upstream LesionDiffusion code: **MIT**. Native package code: **Apache-2.0**, with component-specific notices in the [source distribution](https://github.com/MedMLX/lesiondiffusion-mlx/blob/v0.2.1/THIRD_PARTY_NOTICES.md). **Checkpoint redistribution terms remain unresolved; retrieve original checkpoints directly from the original publisher.**
 
-- `model.safetensors` (1,665,997,956 bytes)
+This entry is a **private converted-weight archive**. Public redistribution review remains separate; the original checkpoint and component terms are unchanged. This update grants no new or broader license and does not clear the model for public redistribution.
 
-Component names prefix tensor keys in `model.safetensors`; the included loader patch selects each component without changing its values or precision.
+Retain the applicable license, copyright and attribution notices, identify the upstream checkpoint and authors, and identify the MLX conversion changes. The existing `NOTICE` and `SOURCE_LICENSE` files are retained without alteration. Review the [original publisher's project](https://github.com/HengruiTianSJTU/LesionDiffusion) and its component terms before use or redistribution.

conversion.json

@@ -10,18 +10,6 @@
     "vq": "float16"
   },
   "ema": true,
-  "source_files": {
-    "LMNet.ckpt": {},
-    "LINet.ckpt": {},
-    "VQ.ckpt": {}
-  },
-  "converted_files": {
-    "lmnet.npz": {},
-    "linet.npz": {},
-    "vq.npz": {},
-    "lmnet_text.npz": {},
-    "linet_text.npz": {}
-  },
   "tensor_file": "model.safetensors",
   "tensor_prefixes": {
     "linet": "linet.",

safetensors-support.patch

@@ -1,130 +0,0 @@
-diff --git a/pyproject.toml b/pyproject.toml
-index bbaab24..e3e0534 100644
---- a/pyproject.toml
-+++ b/pyproject.toml
-@@ -20,6 +20,7 @@ dependencies = [
-   "medmlx-core @ git+https://github.com/MedMLX/medmlx-core.git@0a83421597991ad04c76a532780273c13854cbac",
-   "nibabel>=5.3,<6",
-   "numpy>=2.0,<3",
-+  "safetensors>=0.6,<1",
-   "scipy>=1.15,<2",
-   "tokenizers>=0.21,<1",
- ]
-diff --git a/src/lesiondiffusion_mlx/medmlx.py b/src/lesiondiffusion_mlx/medmlx.py
-index 6dfc21a..e8ec266 100644
---- a/src/lesiondiffusion_mlx/medmlx.py
-+++ b/src/lesiondiffusion_mlx/medmlx.py
-@@ -69,7 +69,7 @@ WARNING: Final = (
- )
- WEIGHTS_HINT: Final = (
-     "Convert explicit upstream checkpoints with scripts/convert_weights.py; "
--    "weights_path must contain the NPZ stages, tokenizer.json and conversion.json."
-+    "weights_path must contain safetensors or NPZ stages, tokenizer.json and conversion.json."
- )
- 
- # Canonical labels match upstream organ_type.json and its inference vocabulary.
-@@ -776,6 +776,12 @@ def source_organ_label(organ: str) -> tuple[int, str]:
- # --- converted weights ----------------------------------------------------------------
- 
- 
-+def _stage_file(directory: Path, name: str) -> Path:
-+    """Use the consolidated release or an existing local NPZ conversion."""
-+    path = directory / "model.safetensors"
-+    return path if path.is_file() else directory / (name + ".npz")
-+
-+
- def require_weights(path: _PathArg, *, inference_mode: str = DEFAULT_INFERENCE_MODE) -> Path:
-     """Require a self-contained, current conversion."""
- 
-@@ -787,17 +793,29 @@ def require_weights(path: _PathArg, *, inference_mode: str = DEFAULT_INFERENCE_M
-             hint=WEIGHTS_HINT,
-         )
-     inference_files = SOURCE_INFERENCE_FILES if inference_mode == "source" else CONVERTED_FILES
--    missing = [
--        directory / name
--        for name in (*inference_files, TOKENIZER_FILE, CONVERSION_FILE)
--        if not (directory / name).is_file()
--    ]
-+    files = [_stage_file(directory, Path(name).stem) for name in inference_files]
-+    files.extend((directory / TOKENIZER_FILE, directory / CONVERSION_FILE))
-+    missing = [file for file in files if not file.is_file()]
-     if missing:
-         raise AssetNotReadyError(
-             f"Native MLX LesionDiffusion weights are not ready: missing {missing[0]}",
-             reason=f"Missing native LesionDiffusion asset: {missing[0]}",
-             hint=WEIGHTS_HINT,
-         )
-+    packed = directory / "model.safetensors"
-+    if packed.is_file():
-+        from safetensors import safe_open
-+
-+        with safe_open(packed, framework="np") as tensors:
-+            names = tensors.keys()
-+            present = {key.split(".", 1)[0] for key in names}
-+        for name in inference_files:
-+            if Path(name).stem not in present:
-+                raise AssetNotReadyError(
-+                    f"Native MLX LesionDiffusion weights are missing stage {Path(name).stem}",
-+                    reason=f"Missing stage in {packed.name}",
-+                    hint=WEIGHTS_HINT,
-+                )
-     try:
-         load_conversion(directory)
-     except (ValueError, OSError) as exc:
-@@ -843,8 +861,15 @@ def load_conversion(directory: Path) -> dict[str, object]:
- def load_stage(directory: Path, name: str, mx: Mlx, *, dtype: Dtype) -> dict[str, Array]:
-     """Load converted arrays and copy them to the MLX device."""
-     load_conversion(directory)
--    path = directory / (name + ".npz")
-+    path = _stage_file(directory, name)
-     loaded = mx.load(str(path))
-+    if path.name == "model.safetensors":
-+        prefix = name + "."
-+        return {
-+            key.removeprefix(prefix): value.astype(dtype)
-+            for key, value in loaded.items()
-+            if key.startswith(prefix)
-+        }
-     return {key: value.astype(dtype) for key, value in loaded.items()}
- 
- 
-diff --git a/uv.lock b/uv.lock
-index 708012c..63dc443 100644
---- a/uv.lock
-+++ b/uv.lock
-@@ -200,6 +200,7 @@ dependencies = [
-     { name = "medmlx-core" },
-     { name = "nibabel" },
-     { name = "numpy" },
-+    { name = "safetensors" },
-     { name = "scipy" },
-     { name = "tokenizers" },
- ]
-@@ -237,6 +238,7 @@ requires-dist = [
-     { name = "pyright", extras = ["nodejs"], marker = "extra == 'dev'", specifier = "==1.1.414" },
-     { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.4,<10" },
-     { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.13,<1" },
-+    { name = "safetensors", specifier = ">=0.6,<1" },
-     { name = "scipy", specifier = ">=1.15,<2" },
-     { name = "scipy-stubs", marker = "extra == 'dev'", specifier = ">=1.18.1.1,<1.19" },
-     { name = "tokenizers", specifier = ">=0.21,<1" },
-diff --git a/typings/safetensors/__init__.pyi b/typings/safetensors/__init__.pyi
-new file mode 100644
-index 0000000..fcfb957
---- /dev/null
-+++ b/typings/safetensors/__init__.pyi
-@@ -0,0 +1,14 @@
-+from os import PathLike
-+from types import TracebackType
-+from typing import Self
-+
-+class safe_open:
-+    def __init__(self, filename: str | PathLike[str], framework: str) -> None: ...
-+    def __enter__(self) -> Self: ...
-+    def __exit__(
-+        self,
-+        exc_type: type[BaseException] | None,
-+        exc_value: BaseException | None,
-+        traceback: TracebackType | None,
-+    ) -> None: ...
-+    def keys(self) -> list[str]: ...