medmlx 0.3.2__py3-none-any.whl

medmlx/__init__.py

@@ -0,0 +1,28 @@
+"""MedMLX model discovery and execution without importing model frameworks."""
+
+from medmlx.api import (
+    describe_model,
+    list_checkpoints,
+    list_models,
+    list_packages,
+    list_setup_plans,
+    model_readiness,
+    run,
+    setup_model,
+    setup_plan,
+)
+from medmlx.probe import runtime_info
+
+__version__ = "0.3.1"
+__all__ = [
+    "describe_model",
+    "list_checkpoints",
+    "list_models",
+    "list_packages",
+    "list_setup_plans",
+    "model_readiness",
+    "run",
+    "runtime_info",
+    "setup_model",
+    "setup_plan",
+]

medmlx/__main__.py

@@ -0,0 +1,7 @@
+"""Entry point for ``python -m medmlx``; behaves like the ``medmlx`` command."""
+
+import sys
+
+from medmlx.cli import main
+
+sys.exit(main())

medmlx/adapters/__init__.py

@@ -0,0 +1 @@
+"""Named model configurations implemented by independently installed runtimes."""

medmlx/adapters/_medvlm_worker.py

@@ -0,0 +1,237 @@
+"""Isolated MedVLM JSON bridge, executed by absolute filename with Python -I.
+
+This file must not import MedMLX or medmlx_core: the selected interpreter owns
+MedVLM's conflicting dependency set. Probe mode reads distribution metadata only.
+Run mode accepts local assets and disables Hub access before importing the model.
+The caller owns timeout, process-group cancellation and final output publication.
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+import os
+import platform
+import sys
+from importlib.metadata import PackageNotFoundError, distribution, version
+from pathlib import Path
+from typing import Protocol, cast
+
+
+class Model(Protocol):
+    config: object
+
+
+class Generation(Protocol):
+    text: object
+    generation_tokens: object
+    prompt_tokens: object
+
+
+class Load(Protocol):
+    def __call__(
+        self, path: str, *, strict: bool, local_files_only: bool,
+    ) -> tuple[Model, object]: ...
+
+
+class Template(Protocol):
+    def __call__(
+        self, processor: object, config: object, prompt: str, *, num_images: int,
+    ) -> object: ...
+
+
+class Generate(Protocol):
+    def __call__(
+        self, model: object, processor: object, prompt: str, *, image: list[str],
+        max_tokens: int, temperature: float, prefill_step_size: int, verbose: bool,
+    ) -> Generation: ...
+
+
+class Utilities(Protocol):
+    load: Load
+
+
+class Prompts(Protocol):
+    apply_chat_template: Template
+
+
+class Generator(Protocol):
+    generate: Generate
+
+
+class Metal(Protocol):
+    def is_available(self) -> bool: ...
+
+
+class Mlx(Protocol):
+    metal: Metal
+    gpu: object
+
+    def set_default_device(self, device: object) -> None: ...
+
+
+def _object(value: object) -> dict[str, object]:
+    if not isinstance(value, dict) or not all(
+        isinstance(key, str) for key in cast(dict[object, object], value)
+    ):
+        raise ValueError("Expected a JSON object with string keys")
+    return cast(dict[str, object], value)
+
+
+def _runtime() -> dict[str, str]:
+    """Check the scoped provider version and its installed runtime dependency graph."""
+    from packaging.requirements import Requirement
+    from packaging.utils import canonicalize_name
+
+    if platform.system() != "Darwin" or platform.machine() != "arm64":
+        raise RuntimeError("MedVLM requires macOS arm64")
+    if sys.version_info[:2] != (3, 13):
+        raise RuntimeError("The isolated bridge currently requires Python 3.13")
+    if version("med-vlm-mlx") != "0.2.0":
+        raise RuntimeError("This provisional bridge targets med-vlm-mlx==0.2.0")
+    try:
+        version("mlx-vlm")
+    except PackageNotFoundError:
+        pass
+    else:
+        raise RuntimeError("Remove the conflicting mlx-vlm distribution from this environment")
+    pending = ["med-vlm-mlx"]
+    found: dict[str, str] = {}
+    while pending:
+        name = canonicalize_name(pending.pop())
+        if name in found:
+            continue
+        package = distribution(name)
+        found[name] = package.version
+        for text in package.requires or []:
+            requirement = Requirement(text)
+            if requirement.marker is not None and not requirement.marker.evaluate({"extra": ""}):
+                continue
+            installed = version(requirement.name)
+            if not requirement.specifier.contains(installed):
+                raise RuntimeError(f"{name} requires {requirement}; installed {installed}")
+            pending.append(requirement.name)
+    return {name: found[name] for name in (
+        "med-vlm-mlx", "mlx", "transformers", "huggingface-hub", "tokenizers",
+    )}
+
+
+def _local_assets(request: dict[str, object]) -> tuple[Path, Path, str, int]:
+    """Check paths and declared architecture without qualifying checkpoint tensors."""
+    model_id = request.get("model_id")
+    if model_id not in ("medgemma", "lingshu"):
+        raise ValueError("model_id must be medgemma or lingshu")
+    model_value, image_value = request.get("model_dir"), request.get("image_path")
+    if not isinstance(model_value, str) or not isinstance(image_value, str):
+        raise ValueError("model_dir and image_path must be local paths")
+    model, image = Path(model_value), Path(image_value)
+    if not model.is_absolute() or not model.is_dir():
+        raise FileNotFoundError("model_dir must be an existing absolute local directory")
+    if not image.is_absolute() or not image.is_file() or image.suffix.lower() not in {
+        ".png", ".jpg", ".jpeg",
+    }:
+        raise FileNotFoundError("image_path must be an existing local PNG or JPEG file")
+    for name in (
+        "config.json", "tokenizer.json", "tokenizer_config.json", "preprocessor_config.json",
+    ):
+        path = model / name
+        if not path.is_file() or path.stat().st_size == 0:
+            raise FileNotFoundError(f"Local checkpoint requires {name}; downloads are disabled")
+    config = _object(json.loads((model / "config.json").read_text()))
+    expected = "gemma3" if model_id == "medgemma" else "qwen2_5_vl"
+    if config.get("model_type") != expected:
+        raise ValueError(f"{model_id} requires model_type={expected}")
+    weights = list(model.glob("*.safetensors"))
+    if not weights or any(path.stat().st_size == 0 for path in weights):
+        raise FileNotFoundError("Local checkpoint requires nonempty safetensors weights")
+    index = model / "model.safetensors.index.json"
+    if index.is_file():
+        mapping = _object(_object(json.loads(index.read_text())).get("weight_map"))
+        if not mapping:
+            raise ValueError("Checkpoint weight_map must not be empty")
+        for shard in mapping.values():
+            if not isinstance(shard, str) or Path(shard).name != shard or not shard.endswith(
+                ".safetensors"
+            ):
+                raise ValueError("Checkpoint shard names must be local safetensors filenames")
+            if not (model / shard).is_file():
+                raise FileNotFoundError(f"Missing checkpoint shard: {shard}")
+    prompt = request.get("prompt")
+    if not isinstance(prompt, str) or not prompt.strip() or len(prompt) > 16384:
+        raise ValueError("prompt must contain 1..16384 characters, including non-whitespace")
+    tokens = request.get("max_new_tokens")
+    if type(tokens) is not int or not 1 <= tokens <= 2048:
+        raise ValueError("max_new_tokens must be an integer in [1, 2048]")
+    return model, image, prompt, tokens
+
+
+def _generate(
+    assets: tuple[Path, Path, str, int], versions: dict[str, str],
+) -> dict[str, object]:
+    """Use the provider's public strict loader, template and greedy generation APIs."""
+    model_path, image_path, prompt, tokens = assets
+    mx = cast(Mlx, importlib.import_module("mlx.core"))
+    if not mx.metal.is_available():
+        raise RuntimeError("The isolated MedVLM interpreter has no MLX Metal device")
+    mx.set_default_device(mx.gpu)
+    load = cast(Utilities, importlib.import_module("mlx_vlm.utils")).load
+    template = cast(Prompts, importlib.import_module("mlx_vlm.prompt_utils")).apply_chat_template
+    generate = cast(Generator, importlib.import_module("mlx_vlm.generate")).generate
+    model, processor = load(str(model_path), strict=True, local_files_only=True)
+    rendered = template(processor, model.config, prompt, num_images=1)
+    if not isinstance(rendered, str):
+        raise TypeError("The medical chat template must return text")
+    result = generate(model, processor, rendered, image=[str(image_path)], max_tokens=tokens,
+                      temperature=0.0, prefill_step_size=64, verbose=False)
+    text: object = result.text
+    generated: object = result.generation_tokens
+    prefilling: object = result.prompt_tokens
+    if not isinstance(text, str) or not text.strip():
+        raise RuntimeError("MedVLM generated no text")
+    if type(generated) is not int or not 1 <= generated <= tokens:
+        raise RuntimeError("MedVLM returned an invalid generation token count")
+    if type(prefilling) is not int or prefilling < 1:
+        raise RuntimeError("MedVLM returned an invalid prompt token count")
+    return {"text": text, "generation_tokens": generated, "prompt_tokens": prefilling,
+            "versions": versions, "python": sys.executable}
+
+
+def main() -> int:
+    """Read one request and atomically return a versioned JSON success/error packet."""
+    if len(sys.argv) != 3:
+        raise SystemExit("usage: _medvlm_worker.py REQUEST_JSON RESPONSE_JSON")
+    request_path, response_path = map(Path, sys.argv[1:])
+    # These are process-local execution settings, not persistent access changes.
+    os.environ.update(HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", HF_DATASETS_OFFLINE="1")
+    stage = "request"
+    try:
+        request = _object(json.loads(request_path.read_text()))
+        if request.get("protocol") != 1:
+            raise ValueError("Unsupported MedVLM bridge protocol")
+        if request.get("action") == "probe":
+            stage = "runtime"
+            result: dict[str, object] = {"versions": _runtime(), "python": sys.executable}
+        elif request.get("action") == "run":
+            stage = "assets"
+            assets = _local_assets(request)
+            stage = "runtime"
+            versions = _runtime()
+            stage = "execution"
+            result = _generate(assets, versions)
+        else:
+            raise ValueError("Unknown MedVLM bridge action")
+        packet: dict[str, object] = {"protocol": 1, "ok": True, "result": result}
+        code = 0
+    except Exception as error:
+        packet = {"protocol": 1, "ok": False, "stage": stage,
+                  "error": {"type": type(error).__name__, "message": str(error)}}
+        code = 1
+    text = json.dumps(packet, allow_nan=False)
+    temporary = response_path.with_suffix(".tmp")
+    temporary.write_text(text)
+    os.replace(temporary, response_path)
+    return code
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

medmlx/adapters/fdg_pet.py

@@ -0,0 +1,295 @@
+"""Explicit-profile, weight-free FDG PET synthesis using fdg-pet-ct-mlx."""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import tempfile
+from collections.abc import Mapping
+from dataclasses import dataclass
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+import numpy as np
+from medmlx_core.errors import InvalidInputError, MissingDependencyError, ModelExecutionError
+from numpy.typing import NDArray
+
+if TYPE_CHECKING:
+    from fdg_pet_ct_mlx._typing import FdgRuntime
+
+MODEL = {
+    "id": "fdg-pet-ct",
+    "label": "Synthetic FDG PET",
+    "task": "generation",
+    "run": "medmlx.adapters.fdg_pet:run",
+    "readiness": "medmlx.adapters.fdg_pet:readiness",
+    "summary": "Weight-free PET SUV synthesis from CT, MAISI labels and explicit profiles.",
+}
+
+
+@dataclass(frozen=True)
+class Tissue:
+    """Caller-provided activity mean and standard deviation in SUV, with MAISI IDs."""
+
+    suv_mean: float
+    suv_std: float
+    maisi_labels: tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class Organs:
+    """Validated tissue profiles, including the fallback for unlisted MAISI labels."""
+
+    tissues: Mapping[str, Tissue]
+
+    def lookup_table(self, size: int) -> tuple[float, ...]:
+        """Map labels below ``size`` to mean SUV; use the fallback for unlisted IDs."""
+        table = [self.tissues["unclassified_soft_tissue"].suv_mean] * size
+        for tissue in self.tissues.values():
+            for label in tissue.maisi_labels:
+                if label < size:
+                    table[label] = tissue.suv_mean
+        return tuple(table)
+
+
+@dataclass(frozen=True)
+class Scanner:
+    """Image-space acquisition parameters; FWHM follows the three array axes in mm."""
+
+    scatter_fraction: float
+    randoms_fraction: float
+    fwhm_mm: tuple[float, float, float]
+    counts_per_suv: float
+
+
+def _object(value: object, name: str, keys: set[str] | None = None) -> dict[str, object]:
+    if not isinstance(value, dict) or not all(
+        isinstance(key, str) for key in cast(dict[object, object], value)
+    ):
+        raise InvalidInputError(f"{name} must be a JSON object")
+    result = cast(dict[str, object], value)
+    if keys is not None and set(result) != keys:
+        raise InvalidInputError(f"{name} must contain exactly: {', '.join(sorted(keys))}")
+    return result
+
+
+def _number(
+    value: object,
+    name: str,
+    *,
+    positive: bool = False,
+    maximum: float | None = None,
+) -> float:
+    if isinstance(value, bool) or not isinstance(value, int | float):
+        raise InvalidInputError(f"{name} must be a finite number")
+    result = float(value)
+    if not math.isfinite(result) or result < 0 or (positive and result == 0):
+        domain = "positive" if positive else "nonnegative"
+        raise InvalidInputError(f"{name} must be finite and {domain}")
+    if maximum is not None and result > maximum:
+        raise InvalidInputError(f"{name} must be at most {maximum}")
+    return result
+
+
+def _profiles(path: Path) -> tuple[Organs, Scanner]:
+    """Parse caller-owned profiles; no tissue statistics or scanner presets are invented."""
+    try:
+        raw: object = json.loads(path.read_text())
+    except (OSError, ValueError) as error:
+        raise InvalidInputError(f"Cannot read profile_path: {error}") from error
+    profile = _object(raw, "profile", {"tissues", "scanner"})
+    tissues: dict[str, Tissue] = {}
+    assigned: set[int] = set()
+    for name, raw_tissue in _object(profile["tissues"], "tissues").items():
+        tissue = _object(raw_tissue, name, {"suv_mean", "suv_std", "maisi_labels"})
+        raw_ids = tissue["maisi_labels"]
+        if not isinstance(raw_ids, list):
+            raise InvalidInputError(f"{name}.maisi_labels must be a list of integers")
+        ids: list[int] = []
+        for label in cast(list[object], raw_ids):
+            if isinstance(label, bool) or not isinstance(label, int) or not 0 <= label <= 200:
+                raise InvalidInputError("Profile MAISI label IDs must be integers in [0, 200]")
+            if label in assigned:
+                raise InvalidInputError(f"MAISI label {label} is assigned more than once")
+            assigned.add(label)
+            ids.append(label)
+        tissues[name] = Tissue(
+            _number(tissue["suv_mean"], f"{name}.suv_mean"),
+            _number(tissue["suv_std"], f"{name}.suv_std"),
+            tuple(ids),
+        )
+    if "unclassified_soft_tissue" not in tissues:
+        raise InvalidInputError("tissues requires unclassified_soft_tissue for unlisted labels")
+    scanner = _object(
+        profile["scanner"],
+        "scanner",
+        {
+            "scatter_fraction",
+            "randoms_fraction",
+            "fwhm_mm",
+            "counts_per_suv",
+        },
+    )
+    raw_fwhm = scanner["fwhm_mm"]
+    if not isinstance(raw_fwhm, list) or len(cast(list[object], raw_fwhm)) != 3:
+        raise InvalidInputError("scanner.fwhm_mm requires three axis FWHM values in mm")
+    fwhm = [_number(value, "scanner.fwhm_mm") for value in cast(list[object], raw_fwhm)]
+    return Organs(tissues), Scanner(
+        _number(scanner["scatter_fraction"], "scatter_fraction", maximum=1),
+        _number(scanner["randoms_fraction"], "randoms_fraction", maximum=1),
+        (fwhm[0], fwhm[1], fwhm[2]),
+        _number(scanner["counts_per_suv"], "counts_per_suv", positive=True),
+    )
+
+
+def run(
+    *,
+    image_path: str,
+    labels_path: str,
+    profile_path: str,
+    output_dir: str,
+    count_level: float = 1.0,
+    seed: int = 0,
+    overwrite: bool = False,
+) -> dict[str, object]:
+    """Synthesize PET SUV on the input CT grid with explicit organ and scanner profiles.
+
+    No learned weights are used. The independent provider owns activity sampling
+    and acquisition approximation. This adapter does not simulate projection
+    physics/reconstruction, resample inputs, add lesions, or establish clinical validity.
+
+    Args:
+        image_path: Finite scalar 3-D NIfTI CT in HU, with spatial units set to mm.
+        labels_path: Co-registered NIfTI MAISI labels, integer IDs in [0, 200],
+            with the CT shape, affine (absolute tolerance 1e-5) and mm units.
+        profile_path: JSON object containing exactly ``tissues`` and ``scanner``.
+            Each tissue contains exactly finite nonnegative ``suv_mean`` and
+            ``suv_std`` (SUV), and a list of globally unique integer ``maisi_labels``
+            in [0, 200]. ``unclassified_soft_tissue`` is required for unlisted IDs.
+            Scanner contains exactly ``scatter_fraction`` and ``randoms_fraction``
+            in [0, 1], ``fwhm_mm`` as three finite nonnegative values in array-axis
+            order, and finite positive ``counts_per_suv``. No profile is built in.
+        output_dir: Write pet.nii.gz as float32 SUV, retaining the CT shape/affine.
+        count_level: Finite positive multiplier on scanner counts_per_suv.
+        seed: Integer in [0, 4294967294]. Activity uses seed; acquisition uses seed+1.
+        overwrite: Replace pet.nii.gz only when true. Input/output aliases are refused.
+
+    Returns:
+        Absolute outputs.pet path and explicit image-space synthesis settings.
+    """
+    count_level = _number(count_level, "count_level", positive=True)
+    if type(seed) is not int or not 0 <= seed <= 4294967294:
+        raise InvalidInputError("seed must be an integer in [0, 4294967294]")
+    if type(overwrite) is not bool:
+        raise InvalidInputError("overwrite must be boolean")
+    image, labels, profile = [
+        Path(item).expanduser().resolve() for item in (image_path, labels_path, profile_path)
+    ]
+    destination = Path(output_dir).expanduser().resolve() / "pet.nii.gz"
+    if destination.resolve() in {image, labels, profile}:
+        raise InvalidInputError("Output must differ from all input paths")
+    if destination.is_dir() or (
+        (destination.exists() or destination.is_symlink()) and not overwrite
+    ):
+        raise InvalidInputError(f"Output exists: {destination}; pass overwrite=True to replace it")
+    organs, scanner = _profiles(profile)
+    if not math.isfinite(count_level * scanner.counts_per_suv):
+        raise InvalidInputError("count_level * counts_per_suv must be finite")
+    try:
+        import nibabel as nib
+    except ImportError as error:
+        raise MissingDependencyError("Install medmlx's generate extra for NIfTI I/O") from error
+    try:
+        ct = nib.load(image)
+        segmentation = nib.load(labels)
+    except (OSError, ValueError) as error:
+        raise InvalidInputError(f"Cannot load CT/labels NIfTI: {error}") from error
+    if not isinstance(ct, nib.Nifti1Image | nib.Nifti2Image) or not isinstance(
+        segmentation, nib.Nifti1Image | nib.Nifti2Image
+    ):
+        raise InvalidInputError("CT and labels must be NIfTI images")
+    if ct.header.get_xyzt_units()[0] != "mm" or segmentation.header.get_xyzt_units()[0] != "mm":
+        raise InvalidInputError("CT and labels must declare NIfTI spatial units in mm")
+    if ct.affine is None or segmentation.affine is None:
+        raise InvalidInputError("CT and labels require an affine")
+    hu = ct.get_fdata(dtype=np.float32)
+    ids = segmentation.get_fdata(dtype=np.float64)
+    affine = ct.affine
+    if hu.ndim != 3 or not hu.size or not np.isfinite(hu).all():
+        raise InvalidInputError("CT must be a nonempty finite scalar 3-D volume in HU")
+    if ids.shape != hu.shape or not np.allclose(affine, segmentation.affine, atol=1e-5, rtol=0):
+        raise InvalidInputError("CT and labels must share shape and affine; no resampling occurs")
+    if (
+        not np.isfinite(ids).all()
+        or (ids < 0).any()
+        or (ids > 200).any()
+        or (ids != np.floor(ids)).any()
+    ):
+        raise InvalidInputError("MAISI labels must be finite integer IDs in [0, 200]")
+    spacing = np.linalg.norm(affine[:3, :3], axis=0)
+    if not np.isfinite(affine).all() or np.linalg.det(affine[:3, :3]) == 0 or (spacing <= 0).any():
+        raise InvalidInputError("CT affine must be finite, nonsingular and have positive spacing")
+    from fdg_pet_ct_mlx.acquisition import apply_acquisition_model
+    from fdg_pet_ct_mlx.activity import synthesize_activity_map
+    from medmlx_core.runtime import import_mlx
+
+    mx = cast("FdgRuntime", import_mlx())
+    activity = synthesize_activity_map(mx.array(ids.astype(np.int32)), organs, seed=seed, mx=mx)
+    pet = apply_acquisition_model(
+        activity,
+        mx.array(hu),
+        scanner=scanner,
+        count_level=count_level,
+        spacing_mm=(float(spacing[0]), float(spacing[1]), float(spacing[2])),
+        seed=seed + 1,
+        mx=mx,
+    )
+    result: NDArray[np.float32] = np.asarray(pet, dtype=np.float32)
+    if result.shape != hu.shape or not np.isfinite(result).all() or (result < 0).any():
+        raise ModelExecutionError("FDG synthesis returned invalid shape or SUV values")
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    header = ct.header.copy()
+    header.set_data_dtype(np.float32)
+    output = type(ct)(result, affine, header)
+    with tempfile.NamedTemporaryFile(
+        dir=destination.parent,
+        suffix=".nii.gz",
+        delete=False,
+    ) as temp:
+        staged = Path(temp.name)
+    try:
+        nib.save(output, staged)
+        if overwrite:
+            os.replace(staged, destination)
+        else:
+            os.link(staged, destination)
+    finally:
+        staged.unlink(missing_ok=True)
+    return {
+        "outputs": {"pet": destination.as_posix()},
+        "details": {
+            "method": "image-space FDG synthesis",
+            "units": "SUV",
+            "seed": seed,
+            "count_level": count_level,
+            "shape": list(result.shape),
+            "profile_path": profile.as_posix(),
+            "requires_human_review": True,
+        },
+    }
+
+
+def readiness() -> dict[str, object]:
+    """Check optional packages; profiles and input contents are validated by run."""
+    try:
+        version("fdg-pet-ct-mlx")
+        version("nibabel")
+    except PackageNotFoundError as error:
+        return {
+            "ready": False,
+            "reason": str(error),
+            "setup_hint": "Install the generate extra: uv sync --locked --extra generate",
+        }
+    return {"ready": True, "reason": None, "setup_hint": None}