From eaf0f0fa00cdcc932b0e61110cdd05973363eb32 Mon Sep 17 00:00:00 2001 From: Paul_Sun <49116282+sun-huangqingbo@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:58:38 -0700 Subject: [PATCH 1/2] ci: add basic package build and installation checks --- .github/workflows/tests.yml | 47 ++++++++++++++ pyproject.toml | 15 ++++- tests/smoke_installed_package.py | 101 +++++++++++++++++++++++++++++++ 3 files changed, 161 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/tests.yml create mode 100644 tests/smoke_installed_package.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..3aeab08 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,47 @@ +name: Build and tests + +on: + push: + branches: [master] + pull_request: + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + build-and-test: + runs-on: ubuntu-latest + timeout-minutes: 20 + env: + CUDA_VISIBLE_DEVICES: "" + HF_HUB_OFFLINE: "1" + MPLBACKEND: Agg + PYTHONDONTWRITEBYTECODE: "1" + steps: + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 + with: + persist-credentials: false + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6 + with: + python-version: "3.11" + cache: pip + cache-dependency-path: pyproject.toml + - name: Install CPU PyTorch + run: python -m pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu + - name: Install package and test tools + run: python -m pip install ".[web,dev]" build + - name: Check dependencies + run: python -m pip check + - name: Run tests + run: python -m unittest discover -s tests -v + - name: Build wheel + run: python -m build --wheel + - name: Install built wheel + run: python -I -m pip install --force-reinstall --no-deps dist/*.whl + - name: Check installed package + run: python -I tests/smoke_installed_package.py -v diff --git a/pyproject.toml b/pyproject.toml index ff97064..7243d1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools>=64", "wheel"] +requires = ["setuptools>=77", "wheel"] build-backend = "setuptools.build_meta" [project] @@ -7,7 +7,8 @@ name = "proticelli" version = "0.1.0" description = "ProtiCelli: Protein Visual Language Model for microscopy image generation" readme = "README.md" -license = {text = "MIT"} +license = "GPL-3.0-only" +license-files = ["LICENSE"] requires-python = ">=3.10,<3.14" authors = [ {name = "Huangqingbo (Paul) Sun", email = "sunh@stanford.edu"}, @@ -32,6 +33,12 @@ dependencies = [ "scikit-image", ] +[project.urls] +Homepage = "https://github.com/CellProfiling/ProtiCelli" +Repository = "https://github.com/CellProfiling/ProtiCelli" +Issues = "https://github.com/CellProfiling/ProtiCelli/issues" +Documentation = "https://github.com/CellProfiling/ProtiCelli#readme" + [project.optional-dependencies] train = [ "tensorboard", @@ -51,8 +58,12 @@ dev = [ [project.scripts] proticelli-web = "proticelli_web.cli:main" +[tool.setuptools] +include-package-data = false + [tool.setuptools.packages.find] include = ["proticelli*", "proticelli_web*"] +namespaces = false [tool.setuptools.package-data] "proticelli.data" = ["*.pkl"] diff --git a/tests/smoke_installed_package.py b/tests/smoke_installed_package.py new file mode 100644 index 0000000..3f59454 --- /dev/null +++ b/tests/smoke_installed_package.py @@ -0,0 +1,101 @@ +"""Run with ``python -I tests/smoke_installed_package.py`` after installation. + +Isolated mode prevents the checkout from satisfying package imports. No learned +weights are downloaded: a tiny real DiT checks dependency and checkpoint APIs. +""" + +from importlib import metadata, resources +import contextlib +import io +from pathlib import Path +import pickle +import sys +import tempfile +import unittest +from unittest.mock import patch + +import numpy as np +import torch + +import proticelli +import proticelli_web +from proticelli import Model +from proticelli.models.dit import DiTTransformer2DModel + + +class InstalledPackageTests(unittest.TestCase): + def test_imports_use_installed_distribution_and_matching_version(self): + checkout = Path(__file__).resolve().parents[1] + self.assertNotEqual(Path(proticelli.__file__).resolve().parent, checkout / "proticelli") + self.assertNotEqual(Path(proticelli_web.__file__).resolve().parent, checkout / "proticelli_web") + distribution = metadata.distribution("proticelli") + self.assertEqual(proticelli.__version__, distribution.version) + self.assertEqual(distribution.metadata["License-Expression"], "GPL-3.0-only") + self.assertIn("LICENSE", " ".join(distribution.metadata.get_all("License-File", []))) + self.assertIn("Repository, https://github.com/CellProfiling/ProtiCelli", + distribution.metadata.get_all("Project-URL", [])) + + def test_bundled_vocabularies_and_gallery_assets_are_usable(self): + data = resources.files("proticelli.data") + wrapper = Model(device="cpu", dtype="float32") + for name, actual in (("antibody_map.pkl", wrapper.protein_map), + ("cell_line_map.pkl", wrapper.cellline_map)): + with self.subTest(name=name): + mapping = pickle.loads((data / name).read_bytes()) + self.assertTrue(mapping) + self.assertEqual(mapping, actual) + self.assertTrue(all(isinstance(k, str) and isinstance(v, int) + for k, v in mapping.items())) + static = resources.files("proticelli_web") / "static" + for name in ("index.html", "app.js", "styles.css", "brand.svg", "proticelli-logo.png"): + with self.subTest(name=name): + self.assertTrue((static / name).read_bytes()) + + def test_real_dit_forward_and_public_checkpoint_roundtrip(self): + with torch.random.fork_rng(devices=[]), tempfile.TemporaryDirectory() as directory: + torch.random.default_generator.manual_seed(73) + network = DiTTransformer2DModel( + num_attention_heads=2, attention_head_dim=4, in_channels=4, + out_channels=2, num_layers=1, sample_size=8, patch_size=2, + num_embeds_ada_norm=16, num_protein_labels=5, num_cell_labels=3, + norm_type="ada_norm_zero_continuous", + ).eval() + image = torch.linspace(-1, 1, 2 * 4 * 8 * 8).reshape(2, 4, 8, 8) + conditions = {"timestep": torch.tensor([0.1, 0.5]), + "protein_labels": torch.tensor([1, 2]), + "cell_line_labels": torch.tensor([1, 2])} + with torch.no_grad(): + before = network(image, **conditions).sample + # Output geometry follows configured channels and input patch grid. + self.assertEqual(before.shape, (2, 2, 8, 8)) + self.assertTrue(torch.isfinite(before).all().item()) + self.assertGreater(before.std().item(), 0) + wrapper = Model(device="cpu", dtype="float32", protein_map={"CUSTOM": 2}, + cellline_map={"CUSTOM_CELL": 1}) + wrapper._model = network + with contextlib.redirect_stdout(io.StringIO()): + wrapper.save(directory) + restored = Model(checkpoint_dir=directory, device="cpu", dtype="float32") + self.assertEqual(restored.protein_map, {"CUSTOM": 2}) + self.assertEqual(restored.cellline_map, {"CUSTOM_CELL": 1}) + for name, value in network.state_dict().items(): + torch.testing.assert_close(value, restored.model.state_dict()[name], rtol=0, atol=0) + with torch.no_grad(): + after = restored.model(image, **conditions).sample + # Independent serialization invariant: preserving every weight and + # configuration must preserve the result for the same real input. + np.testing.assert_array_equal(after.numpy(), before.numpy()) + + def test_installed_console_entry_point_supports_help_without_assets(self): + distribution = metadata.distribution("proticelli") + entry = next(ep for ep in distribution.entry_points if ep.name == "proticelli-web") + output = io.StringIO() + with patch.object(sys, "argv", ["proticelli-web", "--help"]), contextlib.redirect_stdout(output): + with self.assertRaises(SystemExit) as stopped: + entry.load()() + self.assertEqual(stopped.exception.code, 0) + self.assertIn("--mode", output.getvalue()) + + +if __name__ == "__main__": + unittest.main() From 61c3f3248588234cff5f90aafe2ff45004eee7c1 Mon Sep 17 00:00:00 2001 From: Paul_Sun <49116282+sun-huangqingbo@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:00:50 -0700 Subject: [PATCH 2/2] fix: harden runtime input, checkpoint and training edge cases --- README.md | 31 +- proticelli/_training.py | 10 +- proticelli/data/preprocessing.py | 20 +- proticelli/model.py | 28 +- proticelli/models/basic_transformer_block.py | 7 + proticelli/models/dit.py | 32 +- proticelli/utils/checkpoint_utils.py | 2 +- quickstart.ipynb | 63 ++- tests/test_agent_preprocessing.py | 128 +++++ tests/test_core_regressions.py | 512 +++++++++++++++++++ tests/test_input_contract.py | 270 ++++++++++ tests/test_resume_memory.py | 119 +++++ tests/test_training_controls.py | 321 ++++++++++++ tests/test_training_regressions.py | 255 +++++++++ tests/test_transformer_compatibility.py | 143 ++++++ 15 files changed, 1915 insertions(+), 26 deletions(-) create mode 100644 tests/test_agent_preprocessing.py create mode 100644 tests/test_core_regressions.py create mode 100644 tests/test_input_contract.py create mode 100644 tests/test_resume_memory.py create mode 100644 tests/test_training_controls.py create mode 100644 tests/test_training_regressions.py create mode 100644 tests/test_transformer_compatibility.py diff --git a/README.md b/README.md index b817bea..da8bb2d 100644 --- a/README.md +++ b/README.md @@ -118,7 +118,8 @@ python -m proticelli_web --mode real ``` Open . In the default `--mode auto`, the server uses the -real ProtiCelli engine when `proticelli/checkpoint/unet` and `proticelli/vae` +real ProtiCelli engine when either `proticelli/checkpoint/unet` or +`proticelli/checkpoint/unet_ema`, together with `proticelli/vae`, are present. Without those assets, auto mode enters an explicitly labeled **demo renderer** for UI development. `--mode real` instead refuses to start without the assets. The one-click launchers download missing assets @@ -310,6 +311,11 @@ python -m unittest discover -s tests -v ## Installation +The canonical source is [CellProfiling/ProtiCelli](https://github.com/CellProfiling/ProtiCelli). +Package metadata and the bundled license use GPLv3. For development checks, +install `".[web,dev]"` and run `python -m unittest discover -s tests -v`. +GitHub Actions also builds a wheel and checks its installation; it does not publish packages. + ### Python API and training The Gallery is an optional interface; the original Python workflows remain @@ -376,7 +382,7 @@ stack = ChannelAssembler(has_protein=True).transform({ ### 3. Normalize images -The model consumes `[-1, 1]` conditioning. Use `ImageNormalizer` on raw stacks, +The model consumes conditioning on a nominal `[-1, 1]` scale. Use `ImageNormalizer` on raw stacks, then pass `input_range="normalized"` to prediction so an all-positive normalized image is preserved. The four-channel order is always `[microtubules, protein, nucleus, ER]`; the three conditioning channels are @@ -592,7 +598,7 @@ results = model.predict( | `seed` | `int` or `None` | `None` | Random seed for reproducible results. | | `return_latents` | `bool` | `False` | If `True`, includes raw latent tensors in the result object. | | `show_progress` | `bool` | `True` | Show a progress bar during generation. | -| `input_range` | `str` | `"auto"` | `"normalized"` preserves `[-1,1]`, `"unit"` converts `[0,1]`, and `"raw"` uses acquisition bit depth. | +| `input_range` | `str` | `"auto"` | `"normalized"` preserves prepared intensities, including finite overshoot beyond nominal `[-1,1]`; `"unit"` converts `[0,1]`, and `"raw"` uses acquisition bit depth. | | `bit_depth` | `int` or `None` | `None` | `8` or `16` for raw float counts or an explicit acquisition-depth override. | Inputs must be finite, real 512x512 arrays with three or four channels. Resample @@ -609,6 +615,10 @@ Use `input_range="normalized"` whenever the range is known, including positive normalized three-channel arrays. Float counts above one require a declared bit depth, for example `input_range="raw", bit_depth=16`. +Auto mode keeps the pre-existing floating-point convention: arrays wholly in +`[0,1]` are mapped to `[-1,1]`; other finite float arrays pass through as prepared +conditioning. It does not infer an acquisition bit depth from float brightness. + Seeding covers both VAE conditioning samples and denoising noise. Repeatability is defined for the same inputs, batch/solver settings, device, dtype, and software environment; it does not promise identical results across hardware. @@ -850,6 +860,16 @@ Each dict value accepts a file path **or** a numpy array. Files saved as `(1, H, ### `ImageNormalizer` — Normalize to `[-1, 1]` +This is the current preprocessing option for new raw images; it uses the +nucleus anchor and adaptive gain lift described below. It differs from the +[paper's training preprocessing](https://pmc.ncbi.nlm.nih.gov/articles/PMC13060211/), +which clips each channel at its own 99.5th percentile, divides all channels by +the clipped microtubule maximum, and applies `2*x - 1`. That shared denominator +allows other channels to exceed 1 (values around 1.5 are valid). +Pass already prepared training/reference data with `input_range="normalized"` +to preserve it, without applying `ImageNormalizer` again. There is no hard 1.5 +ceiling on finite normalized inputs. + ```python from proticelli.data import ImageNormalizer @@ -975,7 +995,10 @@ Resample **after** normalizing. Interpolation on raw counts changes the percenti **For prediction:** `[512, 512, 3]` arrays ordered `[microtubules, nucleus, ER]`, or `[512, 512, 4]` arrays ordered `[microtubules, protein, nucleus, ER]`. -Use `input_range="normalized"` for known `[-1,1]` floats. The protein channel is +Use `input_range="normalized"` for prepared floats on the nominal `[-1,1]` scale. +Existing reference TIFFs can contain finite values above 1; these are preserved +without clipping or rescaling. Values must remain finite when cast to float32. +The protein channel is excluded from conditioning after validation and range conversion. **For training:** `[H, W, 4]` TIFF where: diff --git a/proticelli/_training.py b/proticelli/_training.py index 2c56268..c322ebd 100644 --- a/proticelli/_training.py +++ b/proticelli/_training.py @@ -26,8 +26,9 @@ class FinetuneDataset(Dataset): Conditioning uses channels 0, 2, 3; channel 1 supplies the protein target. Integer scale comes from acquisition dtype or an explicit bit_depth, never - observed brightness. Specify input_range='normalized' for known [-1,1] - float data, especially if an image happens to have only positive values. + observed brightness. Specify input_range='normalized' for prepared float + data on the nominal [-1,1] scale, especially for all-positive images. + Finite overshoot in previously prepared data is preserved. """ def __init__(self, image_dir: str, image_files: List[str], @@ -176,7 +177,10 @@ def run_finetuning( sigma_min=EDM_CONFIG["SIGMA_MIN"], sigma_max=EDM_CONFIG["SIGMA_MAX"], sigma_data=EDM_CONFIG["SIGMA_DATA"], num_train_timesteps=1000, prediction_type="epsilon") sigmas_all = noise_scheduler.sigmas.to(device=device, dtype=torch.float32) - ema = {name: value.detach().clone() for name, value in dit.state_dict().items()} if use_ema else None + # Resume loads its saved EMA directly; allocating a fresh full shadow + # first would temporarily retain two shadows on the accelerator. + ema = ({name: value.detach().clone() for name, value in dit.state_dict().items()} + if use_ema and state is None else None) epoch_start, next_batch, global_step = 0, 0, 0 if state is not None: dit.load_state_dict(state["model"]) diff --git a/proticelli/data/preprocessing.py b/proticelli/data/preprocessing.py index ed8a065..9a937c3 100644 --- a/proticelli/data/preprocessing.py +++ b/proticelli/data/preprocessing.py @@ -23,13 +23,15 @@ def _normalize_image_range(image, *, input_range="auto", bit_depth=None): - """Convert an explicitly described intensity range to float32 [-1, 1]. + """Convert declared intensity conventions to finite float32 conditioning. Auto uses the acquisition dtype for uint8/uint16. Floating-point arrays - with any negative value are already normalized; nonnegative floats are - unit-range. Inspect the complete four-channel image before discarding the + wholly within [0,1] are unit-range; other floats are already normalized, + preserving the pre-existing behavior. Inspect all channels before discarding the protein placeholder. For ambiguous positive normalized floats, callers must specify input_range='normalized'. Raw float counts require bit_depth. + Normalized data uses a nominal [-1,1] scale; preserve finite overshoot in + already prepared data instead of clipping or normalizing it a second time. """ array = np.asarray(image) if array.ndim != 3 or array.shape[-1] not in (3, 4) or min(array.shape) == 0: @@ -40,6 +42,13 @@ def _normalize_image_range(image, *, input_range="auto", bit_depth=None): raise ValueError("input_range must be auto, normalized, unit, or raw") if bit_depth is not None and bit_depth not in (8, 16): raise ValueError("bit_depth must be 8 or 16") + try: + with np.errstate(over="raise", invalid="raise"): + result = array.astype(np.float32, copy=True) + except (FloatingPointError, OverflowError) as exc: + raise ValueError("Pixel values must be representable as finite float32") from exc + if not np.isfinite(result).all(): + raise ValueError("Pixel values must be representable as finite float32") mode = input_range if mode == "auto": if array.dtype in (np.dtype('uint8'), np.dtype('uint16')): @@ -51,7 +60,7 @@ def _normalize_image_range(image, *, input_range="auto", bit_depth=None): elif bit_depth is not None: mode = "raw" else: - mode = "normalized" if array.min() < 0 else "unit" + mode = "unit" if result.min() >= 0 and result.max() <= 1 else "normalized" if mode == "raw": if bit_depth is None: if array.dtype not in (np.dtype('uint8'), np.dtype('uint16')): @@ -62,12 +71,11 @@ def _normalize_image_range(image, *, input_range="auto", bit_depth=None): else: maximum = 1.0 low, high = (-1.0, 1.0) if mode == "normalized" else (0.0, 1.0) - if array.min() < low or array.max() > high: + if (mode == "raw" or input_range == "unit") and (array.min() < low or array.max() > high): raise ValueError( f"{mode} pixels must lie in [{low}, {high}]. " "Specify input_range and bit_depth explicitly for raw counts." ) - result = array.astype(np.float32, copy=True) if mode == "raw": result /= maximum if mode != "normalized": diff --git a/proticelli/model.py b/proticelli/model.py index 8478b22..14e189a 100644 --- a/proticelli/model.py +++ b/proticelli/model.py @@ -314,9 +314,10 @@ def predict( ignored and channels 0, 2, 3 are used as conditioning. Spatial dimensions must be 512x512 at the model pixel size. input_range : {"auto", "normalized", "unit", "raw"} - Use "normalized" for known [-1,1] floats, "unit" for [0,1], + Use "normalized" for prepared floats on the nominal [-1,1] scale, + preserving finite overshoot; use "unit" for [0,1], or "raw" plus bit_depth for float counts. Auto infers unsigned - integer scale from dtype and treats nonnegative floats as unit. + integer scale from dtype and treats only floats wholly in [0,1] as unit. It inspects all four channels before dropping the protein channel. bit_depth : {8, 16}, optional Acquisition bit depth for raw counts; never inferred from brightness. @@ -446,7 +447,7 @@ def predict( raise FloatingPointError("Decoder returned non-finite predictions") all_images.append(pred_np) if return_latents: - all_latents.append(latents[i].cpu().numpy().copy()) + all_latents.append(latents[i].float().cpu().numpy().copy()) all_meta.append({ "protein_name": protein_names[start + i], "cell_line_name": cell_line_names[start + i], @@ -810,6 +811,12 @@ def save(self, path: Union[str, Path]) -> None: with open(path / "cell_line_map.pkl", "wb") as f: pickle.dump(self.cellline_map, f) + # An existing directory may still contain an older EMA export. Make + # default inference select the current model that was just saved. + with (path / "proticelli_checkpoint.json").open("w", encoding="utf-8") as f: + json.dump({"version": 1, "inference_subfolder": "unet", + "channel_order": ["microtubules", "protein", "nucleus", "ER"]}, f) + print(f"Model saved to {path}") @staticmethod @@ -975,8 +982,19 @@ def _load_model(self, use_ema=False): ) from .models.dit import DiTTransformer2DModel + load_options = {} + config_path = self.checkpoint_dir / subfolder / "config.json" + if config_path.is_file(): + with config_path.open(encoding="utf-8") as handle: + saved_config = json.load(handle) + if (saved_config.get("norm_type") == "ada_norm_zero" + and "legacy_discrete_capacity" not in saved_config + and saved_config.get("num_embeds_ada_norm", 1000) + and saved_config.get("num_embeds_ada_norm", 1000) + < saved_config.get("num_protein_labels", 13349)): + load_options["legacy_discrete_capacity"] = True self._model = DiTTransformer2DModel.from_pretrained( - str(self.checkpoint_dir), subfolder=subfolder + str(self.checkpoint_dir), subfolder=subfolder, **load_options ) self._model.to(self.device, dtype=self.dtype) @@ -1088,7 +1106,7 @@ def _generator(self, seed): return torch.Generator(device=device).manual_seed(int(seed)) def _preprocess_image(self, img: np.ndarray, *, input_range="auto", bit_depth=None) -> torch.Tensor: - """Convert a user-provided numpy image to a [3, H, W] tensor in [-1, 1].""" + """Create [3,H,W] conditioning; preserve already normalized intensities.""" from .data.preprocessing import _normalize_image_range img = _normalize_image_range(img, input_range=input_range, bit_depth=bit_depth) if img.shape[2] == 4: diff --git a/proticelli/models/basic_transformer_block.py b/proticelli/models/basic_transformer_block.py index bf2dc0f..e788706 100644 --- a/proticelli/models/basic_transformer_block.py +++ b/proticelli/models/basic_transformer_block.py @@ -158,6 +158,9 @@ def forward( embedding: Optional[torch.Tensor] = None, cross_attention_kwargs: Optional[Dict[str, Any]] = None, added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + *, + timestep: Optional[torch.Tensor] = None, + class_labels: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Forward pass. @@ -177,6 +180,10 @@ def forward( Extra kwargs for attention processors. added_cond_kwargs : dict, optional Additional conditioning kwargs. + timestep : torch.Tensor, optional + Timestep conditioning for inherited discrete or single AdaLN modes. + class_labels : torch.Tensor, optional + Discrete class conditioning for inherited AdaLN-Zero. Returns ------- diff --git a/proticelli/models/dit.py b/proticelli/models/dit.py index 4aae4e2..0616a25 100644 --- a/proticelli/models/dit.py +++ b/proticelli/models/dit.py @@ -66,6 +66,9 @@ class DiTTransformer2DModel(ModelMixin, ConfigMixin): Size of cell line label vocabulary (including null label at 0). positional_embeddings : torch.Tensor, optional Custom positional embeddings. + legacy_discrete_capacity : bool + Preserve an old discrete AdaLN-Zero class table when loading its weights. + Protein IDs without trained rows are rejected. Defaults to False. """ _skip_layerwise_casting_patterns = ["pos_embed", "norm"] @@ -95,6 +98,7 @@ def __init__( num_protein_labels: int = 13349, num_cell_labels: int = 41, positional_embeddings: Optional[torch.Tensor] = None, + legacy_discrete_capacity: bool = False, ): super().__init__() @@ -102,6 +106,8 @@ def __init__( self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim self.out_channels = in_channels if out_channels is None else out_channels self.gradient_checkpointing = False + if legacy_discrete_capacity and norm_type == "ada_norm_zero" and not num_embeds_ada_norm: + raise ValueError("Legacy discrete checkpoints require their original class capacity") # Dual label embeddings: protein + cell line self.emb_protein_labels = CombinedTimestepLabelEmbeddings( @@ -142,7 +148,12 @@ def __init__( self.config.attention_head_dim, dropout=self.config.dropout, activation_fn=self.config.activation_fn, - num_embeds_ada_norm=self.config.num_embeds_ada_norm, + # Discrete AdaLN-Zero embeds protein class IDs internally. + # Every public protein ID must fit, even when its vocabulary + # exceeds the generic AdaLN embedding-count default. + num_embeds_ada_norm=(max(num_embeds_ada_norm or 0, num_protein_labels) + if norm_type == "ada_norm_zero" and not legacy_discrete_capacity + else num_embeds_ada_norm), attention_bias=self.config.attention_bias, upcast_attention=self.config.upcast_attention, norm_type=norm_type, @@ -202,12 +213,20 @@ def forward( ------- Transformer2DModelOutput or tuple """ + # Preserve old discrete class tables, including the classifier-free + # dropout row, without assigning untrained vectors to additional IDs. + if self.config.norm_type == "ada_norm_zero" and self.config.legacy_discrete_capacity: + capacity = self.config.num_embeds_ada_norm + if torch.any((protein_labels < 0) | (protein_labels >= capacity)).item(): + raise ValueError(f"Legacy discrete checkpoint supports protein labels in [0, {capacity}); " + "higher IDs have no trained class embeddings") + # Embed labels emb_protein = self.emb_protein_labels( - timestep, protein_labels, hidden_dtype=timestep.dtype + timestep, protein_labels, hidden_dtype=hidden_states.dtype ) emb_cell_line = self.emb_cell_line_labels( - timestep, cell_line_labels, hidden_dtype=timestep.dtype + timestep, cell_line_labels, hidden_dtype=hidden_states.dtype ) emb = torch.cat([emb_protein, emb_cell_line], dim=1) @@ -229,9 +248,10 @@ def forward( def create_custom_forward(module, return_dict=None): def custom_forward(*inputs): if return_dict is not None: - return module(*inputs, return_dict=return_dict) + return module(*inputs, return_dict=return_dict, + timestep=timestep, class_labels=protein_labels) else: - return module(*inputs) + return module(*inputs, timestep=timestep, class_labels=protein_labels) return custom_forward ckpt_kwargs = ( @@ -257,6 +277,8 @@ def custom_forward(*inputs): encoder_attention_mask=None, embedding=emb, cross_attention_kwargs=cross_attention_kwargs, + timestep=timestep, + class_labels=protein_labels, ) # Output projection + unpatchify diff --git a/proticelli/utils/checkpoint_utils.py b/proticelli/utils/checkpoint_utils.py index 831535c..f35b1ab 100644 --- a/proticelli/utils/checkpoint_utils.py +++ b/proticelli/utils/checkpoint_utils.py @@ -19,7 +19,7 @@ def save_model_hook(models, weights, output_dir): weights.pop() def load_model_hook(models, input_dir): - from .models.dit import DiTTransformer2DModel + from ..models.dit import DiTTransformer2DModel if args.use_ema and ema_model is not None: try: diff --git a/quickstart.ipynb b/quickstart.ipynb index bd1ab7b..1e9fd1f 100644 --- a/quickstart.ipynb +++ b/quickstart.ipynb @@ -177,6 +177,7 @@ " cell_line_names=[\"A-431\", \"A-431\"],\n", " num_inference_steps=50, # lower (e.g. 20) for a faster preview\n", " batch_size=4,\n", + " input_range=\"normalized\", # preserve prepared reference intensities\n", " seed=42, # remove for non-deterministic sampling\n", ")\n", "\n", @@ -235,7 +236,65 @@ "cell_type": "markdown", "id": "8897d88c", "metadata": {}, - "source": "## Next Steps\n\n**Predict many proteins at once** — pass lists of any length to `model.predict()`;\nuse `batch_size` to control GPU memory usage:\n```python\nproteins = [\"TOMM20\", \"BRD4\", \"LMNA\", \"ACTB\"]\nresults = model.predict(\n images=[ref_image] * len(proteins),\n protein_names=proteins,\n cell_line_names=[\"U2OS\"] * len(proteins),\n batch_size=8,\n)\n```\n\n**Preprocess your own raw images** — see the README's `ImageNormalizer` section\nfor the full algorithm and parameters. In short:\n```python\nfrom proticelli.data import ChannelAssembler, ImageNormalizer, ResolutionResampler\n\nstack = ChannelAssembler(has_protein=False).transform({\n \"microtubules\": \"mt.tif\", \"nucleus\": \"nucleus.tif\", \"er\": \"er.tif\",\n})\nnorm, gains = ImageNormalizer(bit_depth=16).transform(stack, return_gains=True)\nready = ResolutionResampler().transform(norm, xy_resolution=0.0707)\n```\n\n`gains` is the per-channel scale factor `ImageNormalizer` applied (every\nchannel is scaled relative to the nucleus channel, channel 2). Cache it if you\nlater want to map a prediction back onto the input's intensity scale, or to\nsupply an explicit gain for the protein channel — which has no observed\nintensity of its own at inference time, since `ChannelAssembler(has_protein=False)`\nfills it with zeros:\n```python\n# results.images[0] is the model's prediction, [H, W] in [0, 1]\npred_stack = np.zeros_like(stack)\npred_stack[..., 1] = results.images[0] * 255 # embed in the protein channel\nrecon = ImageNormalizer(bit_depth=16).transform(\n pred_stack, gains=gains, clamp_gains=False\n)\n```\n\n**Fine-tune on your own data** — provide a directory of 4-channel TIFFs:\n```python\nmodel.fit(\n image_dir=\"data/train\",\n image_files=[\"cell_0.tiff\", \"cell_1.tiff\"],\n protein_names=[\"CDT1\", \"TUBB\"],\n cell_line_names=[\"U2OS\", \"U2OS\"],\n output_dir=\"finetuned/\",\n num_epochs=50,\n)\n```\n\n**Save and reload a fine-tuned model:**\n```python\nmodel.save(\"finetuned/\")\nmodel = Model(checkpoint_dir=\"finetuned/\")\n```\n" + "source": [ + "## Next Steps\n", + "\n", + "**Predict many proteins at once** — pass lists of any length to `model.predict()`;\n", + "use `batch_size` to control GPU memory usage:\n", + "```python\n", + "proteins = [\"TOMM20\", \"BRD4\", \"LMNA\", \"ACTB\"]\n", + "results = model.predict(\n", + " images=[ref_image] * len(proteins),\n", + " protein_names=proteins,\n", + " cell_line_names=[\"U2OS\"] * len(proteins),\n", + " batch_size=8,\n", + " input_range=\"normalized\",\n", + ")\n", + "```\n", + "\n", + "**Preprocess your own raw images** — see the README's `ImageNormalizer` section\n", + "for the full algorithm and parameters. In short:\n", + "```python\n", + "from proticelli.data import (ChannelAssembler, ImageNormalizer,\n", + " ResolutionResampler, center_crop_or_pad)\n", + "\n", + "stack = ChannelAssembler(has_protein=False).transform({\n", + " \"microtubules\": \"mt.tif\", \"nucleus\": \"nucleus.tif\", \"er\": \"er.tif\",\n", + "})\n", + "norm = ImageNormalizer(bit_depth=16).transform(stack)\n", + "resampled = ResolutionResampler().transform(norm, xy_resolution=0.0707)\n", + "ready = center_crop_or_pad(resampled, size=512)\n", + "results = model.predict([ready], [\"TOMM20\"], [\"U2OS\"], input_range=\"normalized\")\n", + "```\n", + "\n", + "The four channels are `[microtubules, protein, nucleus, ER]`. Physical resampling\n", + "is followed by a center crop or background padding to the native 512 x 512 grid.\n", + "Use `input_range=\"normalized\"` to preserve already normalized values, even when\n", + "all observed reference pixels happen to be positive.\n", + "\n", + "Predictions are model intensities, clipped to `[0, 1.3]`. Reference-only gains\n", + "cannot convert them to calibrated raw protein counts: there is no observed\n", + "protein measurement at inference time, and normalization includes clipping.\n", + "Do not reapply the zero gain from the missing protein channel to predictions.\n", + "\n", + "**Fine-tune on your own data** — provide a directory of 4-channel TIFFs:\n", + "```python\n", + "model.fit(\n", + " image_dir=\"data/train\",\n", + " image_files=[\"cell_0.tiff\", \"cell_1.tiff\"],\n", + " protein_names=[\"CDT1\", \"TUBB\"],\n", + " cell_line_names=[\"U2OS\", \"U2OS\"],\n", + " output_dir=\"finetuned/\",\n", + " num_epochs=50,\n", + ")\n", + "```\n", + "\n", + "**Save and reload a fine-tuned model:**\n", + "```python\n", + "model.save(\"finetuned/\")\n", + "model = Model(checkpoint_dir=\"finetuned/\")\n", + "```\n" + ] } ], "metadata": { @@ -259,4 +318,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/tests/test_agent_preprocessing.py b/tests/test_agent_preprocessing.py new file mode 100644 index 0000000..2d53a93 --- /dev/null +++ b/tests/test_agent_preprocessing.py @@ -0,0 +1,128 @@ +"""File-to-prediction integration at declared pixel size and channel order.""" + +from pathlib import Path +import tempfile +import types +import unittest + +import numpy as np +from tifffile import imwrite + +from proticelli import PredictionResult +from proticelli.agent_tools import run_tool +from proticelli.data import center_crop_or_pad + + +class CapturingPredictor: + def __init__(self): + self.calls = [] + + def predict(self, **kwargs): + self.calls.append(kwargs) + return PredictionResult( + images=[np.zeros((512, 512), np.float32)], + metadata=[{"protein_name": "A", "cell_line_name": "C"}], + ) + + +class AgentPreprocessingTests(unittest.TestCase): + def test_real_web_engine_forwards_the_known_normalized_range(self): + from proticelli_web.inference import RealInferenceEngine + + base = np.linspace(0, 1, 512 * 512, dtype=np.float32).reshape(512, 512) + ensemble = np.stack([base, base * 0.9, np.flipud(base)]) + received = [] + + class ModelBoundary: + def predict_with_reliability(self, **kwargs): + received.append(kwargs) + return types.SimpleNamespace(ensembles=[ensemble], reliability_scores=[0.9]) + + engine = RealInferenceEngine.__new__(RealInferenceEngine) + engine._model = ModelBoundary() + engine.device, engine.dtype = 'cpu', 'float32' + engine.trajectory_batch_size = 1 + result = engine.generate(np.full((512, 512, 3), 0.4, np.float32), + 'A', 'C', 3, 10, 7, lambda *_: None, + input_range='normalized') + self.assertEqual(received[0]['input_range'], 'normalized') + self.assertEqual(len(result.images), 3) + + def test_ema_only_inference_assets_are_ready_without_a_raw_checkpoint(self): + from proticelli.utils.download import assets_ready + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for name in ('checkpoint/unet_ema', 'vae'): + path = root / name + path.mkdir(parents=True) + (path / 'config.json').write_text('{}', encoding='utf-8') + (path / 'diffusion_pytorch_model.safetensors').write_bytes(b'weight-boundary') + self.assertTrue(assets_ready(root)) + + def test_resampling_preserves_physical_scale_and_pads_normalized_background(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for name, counts in (("mt", 70), ("nucleus", 100), ("er", 70)): + imwrite(root / f"{name}.tif", np.full((512, 512), counts, np.uint16)) + predictor = CapturingPredictor() + result = run_tool(predictor, "predict_from_files", { + "microtubules_path": str(root / "mt.tif"), + "nucleus_path": str(root / "nucleus.tif"), "er_path": str(root / "er.tif"), + "protein_names": ["A"], "cell_line_names": ["C"], + "xy_resolution": 0.0707, "bit_depth": 16, + "output_directory": str(root / "output"), + }) + self.assertEqual(result["status"], "ok", result) + self.assertEqual(len(predictor.calls), 1) + supplied = predictor.calls[0] + self.assertEqual(supplied["input_range"], "normalized") + image = supplied["images"][0] + self.assertEqual(image.shape, (512, 512, 4)) + # Physical resampling gives round(512*.0707/.1067)=339 pixels. + # Centering173extra pixels gives86before and87after each axis. + np.testing.assert_array_equal(image[:86], -1) + np.testing.assert_array_equal(image[425:], -1) + np.testing.assert_array_equal(image[:, :86], -1) + np.testing.assert_array_equal(image[:, 425:], -1) + expected = np.broadcast_to(np.array([0.4, -1, 1, 0.4]), (339, 339, 4)) + np.testing.assert_allclose(image[86:425, 86:425], expected, atol=1e-6) + geometry = result["preprocessing"]["geometry"] + self.assertEqual(geometry["source_shape"], [339, 339]) + self.assertEqual(geometry["padding_tblr"], [86, 87, 86, 87]) + self.assertTrue(all(Path(name).is_file() for name in result["saved_files"])) + + def test_center_crop_preserves_pixels_without_resizing(self): + source = np.arange(5 * 7 * 3, dtype=np.float32).reshape(5, 7, 3) + cropped, geometry = center_crop_or_pad(source, size=3, return_metadata=True) + np.testing.assert_array_equal(cropped, source[1:4, 2:5]) + self.assertEqual(geometry["crop_origin_yx"], [1, 2]) + self.assertEqual(geometry["padding_tblr"], [0, 0, 0, 0]) + + def test_center_padding_places_unequal_margins_consistently(self): + source = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4) + result, geometry = center_crop_or_pad(source, size=6, return_metadata=True) + expected = np.full((6, 6, 4), -1, np.float32) + expected[2:4, 1:4] = source + np.testing.assert_array_equal(result, expected) + self.assertEqual(geometry["padding_tblr"], [2, 2, 1, 2]) + + def test_invalid_crop_inputs_are_rejected(self): + for source in (np.empty((0, 3, 4)), np.full((2, 3, 4), np.nan)): + with self.subTest(shape=source.shape): + with self.assertRaises(ValueError): + center_crop_or_pad(source) + + def test_web_auto_bit_depth_uses_uint16_dtype_before_optional_resampling(self): + from proticelli_web.imaging import prepare_reference_channels + + channels = {"microtubules": np.full((512, 512), 70, np.uint16), + "nucleus": np.full((512, 512), 100, np.uint16), + "er": np.full((512, 512), 70, np.uint16)} + _, metadata = prepare_reference_channels(channels, normalize=True, + pixel_size_um=0.1067, resample=False) + self.assertEqual(metadata["normalization"]["bit_depth"], 16) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_core_regressions.py b/tests/test_core_regressions.py new file mode 100644 index 0000000..22cb4be --- /dev/null +++ b/tests/test_core_regressions.py @@ -0,0 +1,512 @@ +"""Tests-first baseline for the public ProtiCelli inference API. + +These tests intentionally expose existing defects. They do not mark failures as +expected, replace core behavior with fixes, load learned weights, or use a GPU. +Only the heavyweight checkpoint, VAE, denoiser, and prediction-sampler boundaries +are replaced. Numerical sampling tests execute the real EDM sampling functions. +""" + +from __future__ import annotations + +import contextlib +import io +import json +from pathlib import Path +import sys +import tempfile +import types +import unittest +from unittest.mock import patch +import weakref + +import numpy as np +import torch + +from proticelli import Model +from proticelli._sampling import sample_edm, sample_edm_uncertainty +from proticelli.data import ChannelAssembler, ImageNormalizer + + +class TinyCheckpoint(torch.nn.Module): + """A serializable learned-weight boundary with actual Torch state/dtypes.""" + + def __init__(self, value): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([value], dtype=torch.float32)) + self.config = {} + + def register_to_config(self, **options): + self.config.update(options) + + @classmethod + def from_pretrained(cls, directory, *, subfolder): + values = np.load(Path(directory) / subfolder / "boundary_weights.npy") + return cls(float(values[0])) + + def save_pretrained(self, directory): + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + np.save(directory / "boundary_weights.npy", self.weight.detach().cpu().numpy()) + (directory / "config.json").write_text("{}", encoding="utf-8") + + +@contextlib.contextmanager +def checkpoint_boundary(): + """Keep unsupported diffusers/transformers imports outside these unit tests.""" + package = types.ModuleType("proticelli.models") + package.__path__ = [] + module = types.ModuleType("proticelli.models.dit") + module.DiTTransformer2DModel = TinyCheckpoint + # An accidental random initialization must remain observable as weight zero. + module.create_dit_model = lambda **_kwargs: TinyCheckpoint(0) + replacements = {"proticelli.models": package, "proticelli.models.dit": module} + absent = object() + previous = {name: sys.modules.get(name, absent) for name in replacements} + sys.modules.update(replacements) + try: + yield + finally: + # Restore only the checkpoint boundary modules. Restoring the entire + # module registry would discard lazy Torch registrations loaded meanwhile. + for name, original in previous.items(): + if original is absent: + sys.modules.pop(name, None) + else: + sys.modules[name] = original + + +class CheckpointRegressionTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory(dir=Path.cwd()) + self.addCleanup(self.temporary.cleanup) + self.directory = Path(self.temporary.name) + + def model(self, checkpoint=None, **kwargs): + return Model( + checkpoint_dir=self.directory if checkpoint is None else checkpoint, + device="cpu", **kwargs, + ) + + def test_raw_unet_checkpoint_loads_known_weights(self): + TinyCheckpoint(3.25).save_pretrained(self.directory / "unet") + with checkpoint_boundary(): + loaded = self.model().model + self.assertEqual(loaded.weight.item(), 3.25) + self.assertFalse(loaded.training) + self.assertFalse(loaded.weight.requires_grad) + + def test_explicit_ema_load_selects_ema_over_raw_weights(self): + TinyCheckpoint(3.25).save_pretrained(self.directory / "unet") + TinyCheckpoint(7.5).save_pretrained(self.directory / "unet_ema") + instance = self.model() + with checkpoint_boundary(): + instance._load_model(use_ema=True) + self.assertEqual(instance.model.weight.item(), 7.5) + + def test_ema_only_checkpoint_is_loaded_for_public_inference(self): + TinyCheckpoint(7.5).save_pretrained(self.directory / "unet_ema") + with checkpoint_boundary(): + loaded = self.model().model + self.assertEqual( + loaded.weight.item(), 7.5, + "An EMA-only trained checkpoint must not become random initialization", + ) + + def test_missing_checkpoint_cannot_silently_create_an_untrained_predictor(self): + with checkpoint_boundary(): + with self.assertRaises((FileNotFoundError, ValueError, RuntimeError)): + _ = self.model(self.directory / "missing").model + + def test_existing_unet_directory_without_weights_fails(self): + (self.directory / "unet").mkdir() + with checkpoint_boundary(): + with self.assertRaises((FileNotFoundError, ValueError, RuntimeError)): + _ = self.model().model + + def test_saved_weights_reload_through_the_public_constructor(self): + instance = self.model(protein_map={"CUSTOM": 37}, cellline_map={"CUSTOM_CELL": 9}) + instance._model = TinyCheckpoint(6.75) + destination = self.directory / "saved" + with contextlib.redirect_stdout(io.StringIO()): + instance.save(destination) + with checkpoint_boundary(): + reloaded = self.model(destination).model + self.assertEqual(reloaded.weight.item(), 6.75) + + def test_saving_over_an_ema_checkpoint_reloads_the_new_weights(self): + TinyCheckpoint(7.5).save_pretrained(self.directory / "unet_ema") + (self.directory / "proticelli_checkpoint.json").write_text( + json.dumps({"version": 1, "inference_subfolder": "unet_ema"}), encoding="utf-8" + ) + instance = self.model(protein_map={"CUSTOM": 1}, cellline_map={"C": 1}) + instance._model = TinyCheckpoint(3.25) + with contextlib.redirect_stdout(io.StringIO()): + instance.save(self.directory) + with checkpoint_boundary(): + reloaded = self.model().model + self.assertEqual(reloaded.weight.item(), 3.25) + + def test_registered_accelerate_load_hook_restores_saved_weights(self): + from proticelli.utils.checkpoint_utils import setup_checkpoint_hooks + + TinyCheckpoint(4.5).save_pretrained(self.directory / "unet") + target = TinyCheckpoint(0) + hooks = {} + accelerator = types.SimpleNamespace( + is_main_process=True, device="cpu", + register_save_state_pre_hook=lambda hook: hooks.update(save=hook), + register_load_state_pre_hook=lambda hook: hooks.update(load=hook), + ) + setup_checkpoint_hooks(accelerator, types.SimpleNamespace(use_ema=False)) + models = [target] + with checkpoint_boundary(): + hooks["load"](models, str(self.directory)) + self.assertEqual(models, []) + self.assertEqual(target.weight.item(), 4.5) + + def test_saved_custom_label_maps_override_bundled_defaults_when_reloaded(self): + instance = self.model(protein_map={"CUSTOM": 37}, cellline_map={"CUSTOM_CELL": 9}) + instance._model = TinyCheckpoint(6.75) + destination = self.directory / "saved" + with contextlib.redirect_stdout(io.StringIO()): + instance.save(destination) + reloaded = self.model(destination) + self.assertEqual( + (reloaded.protein_map.get("CUSTOM"), reloaded.cellline_map.get("CUSTOM_CELL")), + (37, 9), + "Model.save() label maps must accompany the saved model on reload", + ) + self.assertEqual(set(reloaded.protein_map), {"CUSTOM"}) + self.assertEqual(set(reloaded.cellline_map), {"CUSTOM_CELL"}) + + +class SigmaSchedule: + def __init__(self, sigmas): + self._values = tuple(sigmas) + + def set_timesteps(self, count): + if count != len(self._values) - 1: + raise AssertionError("The scalar reference schedule has a fixed step count") + self.sigmas = torch.tensor(self._values, dtype=torch.float64) + + +class ZeroRawDenoiser: + def __call__(self, image, timestep, **_kwargs): + return types.SimpleNamespace(sample=torch.zeros_like(image[:, :16])) + + +def scalar_zero_denoiser_reference(initial, sigmas, solver): + """Integrate dx/dsigma = sigma*x/(sigma**2 + .25) with scalar arithmetic. + + A zero raw EDM network gives D(x,sigma)=.25*x/(sigma**2+.25). + The terminal transition to zero uses Euler: denoiser conditioning log(sigma) + is undefined at zero, so a second denoiser evaluation is not permitted there. + """ + value = float(initial) * sigmas[0] + for sigma, next_sigma in zip(sigmas[:-1], sigmas[1:]): + derivative = sigma * value / (sigma * sigma + 0.25) + step = next_sigma - sigma + predicted = value + step * derivative + if solver == "heun" and next_sigma > 0: + next_derivative = next_sigma * predicted / (next_sigma**2 + 0.25) + value += step * (derivative + next_derivative) / 2 + else: + value = predicted + return value + + +class SamplingRegressionTests(unittest.TestCase): + def uncertainty(self, sigmas, *, solver="euler", **kwargs): + return sample_edm_uncertainty( + ZeroRawDenoiser(), SigmaSchedule(sigmas), batch_size=1, image_size=2, + num_inference_steps=len(sigmas) - 1, device="cpu", weight_dtype=torch.float64, + initial_latents=torch.ones((1, 16, 2, 2), dtype=torch.float64), + solver=solver, **kwargs, + ) + + def test_plain_edm_euler_matches_independent_scalar_ode_reference(self): + sigmas = (1.0, 0.5, 0.0) + seed = 17 + initial = torch.randn( + (2, 16, 2, 2), generator=torch.Generator().manual_seed(seed), dtype=torch.float64 + ) + actual = sample_edm( + ZeroRawDenoiser(), SigmaSchedule(sigmas), batch_size=2, image_size=2, + num_inference_steps=2, device="cpu", weight_dtype=torch.float64, + generator=torch.Generator().manual_seed(seed), + ) + expected_factor = scalar_zero_denoiser_reference(1.0, sigmas, "euler") + self.assertAlmostEqual(expected_factor, 0.3) + torch.testing.assert_close(actual, initial * expected_factor, rtol=1e-12, atol=1e-12) + + def test_uncertainty_euler_matches_independent_scalar_ode_reference(self): + sigmas = (1.0, 0.5, 0.25, 0.0) + actual = self.uncertainty(sigmas) + expected = scalar_zero_denoiser_reference(1.0, sigmas, "euler") + torch.testing.assert_close(actual, torch.full_like(actual, expected), rtol=1e-12, atol=1e-12) + + def test_heun_positive_endpoint_matches_independent_scalar_reference(self): + sigmas = (1.0, 0.5, 0.25) + actual = self.uncertainty(sigmas, solver="heun") + expected = scalar_zero_denoiser_reference(1.0, sigmas, "heun") + self.assertAlmostEqual(expected, 0.52) + torch.testing.assert_close(actual, torch.full_like(actual, expected), rtol=1e-12, atol=1e-12) + + def test_heun_terminal_zero_uses_the_final_euler_transition(self): + sigmas = (1.0, 0.5, 0.0) + actual = self.uncertainty(sigmas, solver="heun") + expected = scalar_zero_denoiser_reference(1.0, sigmas, "heun") + self.assertAlmostEqual(expected, 0.325) + self.assertTrue(torch.isfinite(actual).all().item(), "Terminal sigma zero produced non-finite latents") + torch.testing.assert_close(actual, torch.full_like(actual, expected), rtol=1e-12, atol=1e-12) + + def test_single_heun_step_to_zero_remains_finite(self): + actual = self.uncertainty((1.0, 0.0), solver="heun") + self.assertTrue(torch.isfinite(actual).all().item()) + torch.testing.assert_close(actual, torch.full_like(actual, 0.2), rtol=1e-12, atol=1e-12) + + def test_sampler_cancellation_propagates_before_denoising(self): + with self.assertRaises(InterruptedError): + self.uncertainty((1.0, 0.5, 0.0), cancel_check=lambda: True) + + +def bare_model(): + return Model(device="cpu", protein_map={"A": 1}, cellline_map={"C": 1}) + + +class PreprocessingRegressionTests(unittest.TestCase): + def test_raw_uint8_channels_use_the_full_dtype_range(self): + image = np.asarray([[[0, 128, 255], [255, 64, 0]]], dtype=np.uint8) + expected = np.moveaxis(image.astype(np.float32) / 255 * 2 - 1, -1, 0) + np.testing.assert_allclose(bare_model()._preprocess_image(image).numpy(), expected) + + def test_dim_raw_uint16_channels_do_not_infer_eight_bit_depth(self): + image = np.asarray([[[0, 128, 255], [32, 64, 200]]], dtype=np.uint16) + expected = np.moveaxis(image.astype(np.float32) / 65535 * 2 - 1, -1, 0) + np.testing.assert_allclose(bare_model()._preprocess_image(image).numpy(), expected) + + def test_signed_normalized_float_channels_are_preserved(self): + image = np.asarray([[[-0.8, 0.2, 0.9], [0.1, -0.3, 0.7]]], dtype=np.float32) + np.testing.assert_array_equal( + bare_model()._preprocess_image(image).numpy(), np.moveaxis(image, -1, 0) + ) + + def test_four_channel_preprocessing_ignores_the_target_channel(self): + image = np.asarray([[[0, 255, 128, 255], [255, 0, 64, 32]]], dtype=np.uint8) + expected = np.moveaxis(image[..., [0, 2, 3]].astype(np.float32) / 255 * 2 - 1, -1, 0) + actual = bare_model()._preprocess_image(image).numpy() + np.testing.assert_allclose(actual, expected) + + def test_normalizer_pipeline_preserves_conditioning_with_negative_pixels(self): + ramp = np.asarray([[0, 20], [40, 100]], dtype=np.uint8) + stack = ChannelAssembler(has_protein=False).transform({ + "microtubules": ramp, "nucleus": ramp, "er": ramp, + }) + normalized = ImageNormalizer(bit_depth=8).transform(stack) + expected = np.moveaxis(normalized[..., [0, 2, 3]], -1, 0) + np.testing.assert_array_equal(bare_model()._preprocess_image(normalized).numpy(), expected) + + def test_normalizer_pipeline_preserves_conditioning_when_all_reference_pixels_are_positive(self): + stack = ChannelAssembler(has_protein=False).transform({ + "microtubules": np.full((2, 2), 70, dtype=np.uint8), + "nucleus": np.full((2, 2), 100, dtype=np.uint8), + "er": np.full((2, 2), 70, dtype=np.uint8), + }) + normalized = ImageNormalizer(bit_depth=8).transform(stack) + # These are known signed-normalized values, not raw [0,1] inputs: + # 2*(70/100)-1=.4 for MT/ER, and 1 for the nucleus reference. + expected = np.broadcast_to(np.asarray([0.4, 1, 0.4], dtype=np.float32)[:, None, None], (3, 2, 2)) + np.testing.assert_allclose(bare_model()._preprocess_image(normalized).numpy(), expected, atol=1e-7) + + +class LatentDistribution: + def __init__(self, mean, stochastic): + self.mean = mean + self.stochastic = stochastic + + def sample(self, generator=None): + if not self.stochastic: + return self.mean.clone() + return self.mean + torch.randn( + self.mean.shape, generator=generator, device=self.mean.device, dtype=self.mean.dtype + ) + + def mode(self): + return self.mean.clone() + + +class SmallVAE: + config = types.SimpleNamespace(scaling_factor=4) + + def __init__(self, *, stochastic=False, before_encode=None): + self.stochastic = stochastic + self.before_encode = before_encode + self.encoded_batches = [] + + def encode(self, condition): + if self.before_encode is not None: + self.before_encode() + self.encoded_batches.append(len(condition)) + downsampled = torch.nn.functional.avg_pool2d(condition, kernel_size=8) + mean = downsampled.mean(dim=1, keepdim=True).repeat(1, 16, 1, 1) + return types.SimpleNamespace(latent_dist=LatentDistribution(mean, self.stochastic)) + + +def conditioning_image(index=0): + ramp = np.linspace(-0.5, 0.5, 512 * 512, dtype=np.float32).reshape(512, 512) + return np.stack((ramp, ramp + 0.05, ramp - 0.05), axis=-1) + index * 0.01 + + +def boundary_sampler(**kwargs): + """Deterministic conditional response given reference latents and noise.""" + reference = kwargs["reference_channels"] + initial = kwargs.get("initial_latents") + if initial is None: + noise = torch.randn(reference.shape, generator=kwargs.get("generator"), dtype=reference.dtype) + else: + noise = initial[..., :reference.shape[-2], :reference.shape[-1]] + return reference + 0.1 * noise + + +@contextlib.contextmanager +def prediction_boundaries(*, stochastic=False, before_encode=None): + instance = bare_model() + instance._vae = SmallVAE(stochastic=stochastic, before_encode=before_encode) + instance._model = object() + instance._scheduler = object() + instance._decode_latents = lambda latents: torch.nn.functional.interpolate( + latents[:, :1].detach(), scale_factor=8, mode="nearest" + ) + with patch("proticelli._sampling.sample_edm", side_effect=boundary_sampler), patch( + "proticelli._sampling.sample_edm_uncertainty", side_effect=boundary_sampler + ): + yield instance + + +class PredictionOrchestrationRegressionTests(unittest.TestCase): + def setUp(self): + state = torch.random.get_rng_state() + self.addCleanup(torch.random.set_rng_state, state) + torch.random.default_generator.manual_seed(193) + + def predict(self, instance, **kwargs): + return instance.predict([conditioning_image()], ["A"], ["C"], show_progress=False, **kwargs) + + def reliability(self, instance, **kwargs): + options = {"num_samples": 5, "batch_size": 2, "seed": 23, "show_progress": False} + options.update(kwargs) + return instance.predict_with_reliability([conditioning_image()], ["A"], ["C"], **options) + + def test_seeded_plain_prediction_control_with_deterministic_encoder(self): + with prediction_boundaries() as instance: + first = self.predict(instance, seed=23) + second = self.predict(instance, seed=23) + np.testing.assert_array_equal(first.images[0], second.images[0]) + + def test_different_plain_prediction_seeds_change_noise(self): + with prediction_boundaries() as instance: + first = self.predict(instance, seed=23) + second = self.predict(instance, seed=24) + self.assertFalse(np.array_equal(first.images[0], second.images[0])) + + def test_bfloat16_latent_export_preserves_representable_values_in_numpy(self): + # A known latent at the learned sampler boundary isolates export from + # the diffusion algorithm. 0.3125 is exactly representable in BF16. + latent = torch.full((1, 16, 64, 64), 0.3125, dtype=torch.bfloat16) + with prediction_boundaries() as instance: + instance.dtype = torch.bfloat16 + instance._decode_latents = lambda value: torch.nn.functional.interpolate( + value[:, :1].float(), scale_factor=8, mode="nearest" + ) + with patch("proticelli._sampling.sample_edm", return_value=latent): + result = self.predict(instance, return_latents=True) + self.assertEqual(result.latents[0].dtype, np.float32) + np.testing.assert_array_equal(result.latents[0], np.full((16, 64, 64), 0.3125)) + + def test_seeded_prediction_includes_stochastic_conditioning_encode(self): + with prediction_boundaries(stochastic=True) as instance: + first = self.predict(instance, seed=23) + second = self.predict(instance, seed=23) + np.testing.assert_array_equal(first.images[0], second.images[0]) + + def test_seeded_reliability_includes_stochastic_conditioning_encode(self): + with prediction_boundaries(stochastic=True) as instance: + first = self.reliability(instance, return_ensembles=True) + second = self.reliability(instance, return_ensembles=True) + np.testing.assert_array_equal(first.ensembles[0], second.ensembles[0]) + np.testing.assert_array_equal(first.images[0], second.images[0]) + np.testing.assert_array_equal(first.reliability_scores, second.reliability_scores) + + def test_requested_ensemble_medoid_and_score_match_independent_correlations(self): + with prediction_boundaries() as instance: + result = self.reliability(instance, return_ensembles=True) + ensemble = result.ensembles[0] + correlation = np.corrcoef(ensemble.reshape(len(ensemble), -1).astype(np.float64)) + centrality = (correlation.sum(axis=1) - 1) / (len(ensemble) - 1) + medoid = int(np.argmax(centrality)) + score = np.clip((correlation.sum() - len(ensemble)) / (len(ensemble) * (len(ensemble) - 1)), 0, 1) + np.testing.assert_array_equal(result.images[0], ensemble[medoid]) + self.assertAlmostEqual(result.reliability_scores[0], score, places=6) + + def test_unrequested_ensemble_storage_is_not_retained_by_the_medoid(self): + with prediction_boundaries() as instance: + result = self.reliability(instance, return_ensembles=False) + self.assertIsNone(result.ensembles) + retained_roots = {} + for image in result.images: + root = image + while isinstance(root.base, np.ndarray): + root = root.base + retained_roots[id(root)] = root + retained_bytes = sum(root.nbytes for root in retained_roots.values()) + output_bytes = sum(image.nbytes for image in result.images) + self.assertLessEqual( + retained_bytes, output_bytes, + "A returned medoid retains backing storage for unrequested ensemble members", + ) + + def test_encoder_batches_follow_requested_limit_and_preserve_output_count(self): + images = [conditioning_image(index) for index in range(7)] + with prediction_boundaries() as instance: + result = instance.predict(images, ["A"] * 7, ["C"] * 7, + batch_size=2, seed=23, show_progress=False) + batches = instance.vae.encoded_batches + self.assertEqual(batches, [2, 2, 2, 1]) + self.assertEqual(len(result.images), 7) + + def assert_bounded_preprocessing(self, reliability): + references = [] + batch_size = 2 + + def before_encode(): + live = sum(reference() is not None for reference in references) + # Allow an active and previous batch. Retaining all eight inputs + # violates a bound proportional to batch_size, even on tiny images. + self.assertLessEqual(live, 2 * batch_size, "All preprocessed inputs were retained at once") + + with prediction_boundaries(before_encode=before_encode) as instance: + original = instance._preprocess_image + + def monitored(image, **options): + tensor = original(image, **options) + references.append(weakref.ref(tensor)) + return tensor + + with patch.object(instance, "_preprocess_image", side_effect=monitored): + method = instance.predict_with_reliability if reliability else instance.predict + kwargs = {"num_samples": 2} if reliability else {} + result = method( + [conditioning_image(index) for index in range(8)], ["A"] * 8, ["C"] * 8, + batch_size=batch_size, seed=23, show_progress=False, **kwargs, + ) + self.assertEqual(len(result.images), 8) + + def test_plain_prediction_preprocesses_a_bounded_number_of_images(self): + self.assert_bounded_preprocessing(reliability=False) + + def test_reliability_preprocesses_a_bounded_number_of_images(self): + self.assert_bounded_preprocessing(reliability=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_input_contract.py b/tests/test_input_contract.py new file mode 100644 index 0000000..61836ef --- /dev/null +++ b/tests/test_input_contract.py @@ -0,0 +1,270 @@ +"""Normalization and dtype contract, written before production changes. + +Expected values follow declared acquisition bit depth or the explicitly chosen +intensity range. Auto mode uses array dtype and the whole four-channel stack, +including a signed normalized protein placeholder, before selecting references. +""" + +from __future__ import annotations + +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + +import numpy as np +from tifffile import imwrite + +from proticelli import Model +from proticelli._training import FinetuneDataset +from proticelli.data import ChannelAssembler, ImageNormalizer + + +def small_model(): + return Model(device="cpu", protein_map={"A": 1}, cellline_map={"C": 1}) + + +class InputRangeTests(unittest.TestCase): + def test_auto_uint8_one_count_is_scaled_by_dtype(self): + image = np.broadcast_to(np.asarray([0, 1, 1], dtype=np.uint8), (2, 3, 3)) + expected = np.broadcast_to(np.asarray([-1, 2 / 255 - 1, 2 / 255 - 1])[:, None, None], (3, 2, 3)) + np.testing.assert_allclose(small_model()._preprocess_image(image).numpy(), expected, atol=1e-7) + + def test_auto_uint16_dim_values_keep_sixteen_bit_scale(self): + image = np.broadcast_to(np.asarray([20, 128, 255], dtype=np.uint16), (2, 3, 3)) + expected = np.moveaxis(image.astype(np.float64) / 65535 * 2 - 1, -1, 0) + np.testing.assert_allclose(small_model()._preprocess_image(image).numpy(), expected, atol=1e-7) + + def test_auto_preserves_actual_normalizer_pipeline_before_dropping_placeholder(self): + stack = ChannelAssembler(has_protein=False).transform({ + "microtubules": np.full((2, 3), 70, np.uint8), + "nucleus": np.full((2, 3), 100, np.uint8), + "er": np.full((2, 3), 70, np.uint8), + }) + normalized = ImageNormalizer(bit_depth=8).transform(stack) + # Normalizer's signed placeholder establishes the entire stack's range. + self.assertEqual(normalized[0, 0, 1], -1) + expected = np.broadcast_to(np.asarray([0.4, 1, 0.4])[:, None, None], (3, 2, 3)) + np.testing.assert_allclose(small_model()._preprocess_image(normalized).numpy(), expected, atol=1e-7) + + def test_explicit_normalized_positive_floats_are_preserved(self): + image = np.broadcast_to(np.asarray([0.2, 0.4, 0.8], np.float32), (2, 3, 3)) + actual = small_model()._preprocess_image(image, input_range="normalized") + np.testing.assert_array_equal(actual.numpy(), np.moveaxis(image, -1, 0)) + + def test_explicit_normalized_overshoot_is_preserved_without_clipping(self): + # A shared training normalization anchor does not bound every channel + # at 1. Already normalized intensities must reach the encoder as-is. + image = np.broadcast_to(np.asarray([0.2, 1.2511415, 1.5], np.float32), (2, 3, 3)) + actual = small_model()._preprocess_image(image, input_range="normalized") + np.testing.assert_array_equal(actual.numpy(), np.moveaxis(image, -1, 0)) + + def test_auto_signed_normalized_overshoot_is_preserved(self): + image = np.broadcast_to(np.asarray([-1.02, 1.0, 1.5], np.float32), (2, 3, 3)) + np.testing.assert_array_equal(small_model()._preprocess_image(image).numpy(), + np.moveaxis(image, -1, 0)) + + def test_normalized_float64_values_must_remain_finite_after_float32_cast(self): + with self.assertRaises(ValueError): + small_model()._preprocess_image(np.full((2, 3, 3), 1e40, np.float64), + input_range="normalized") + + def test_explicit_unit_floats_are_scaled_once(self): + image = np.broadcast_to(np.asarray([0.0, 0.25, 1.0], np.float32), (2, 3, 3)) + actual = small_model()._preprocess_image(image, input_range="unit") + expected = np.broadcast_to(np.asarray([-1, -0.5, 1])[:, None, None], (3, 2, 3)) + np.testing.assert_array_equal(actual.numpy(), expected) + + def test_auto_positive_unit_floats_keep_legacy_unit_interpretation(self): + image = np.broadcast_to(np.asarray([0.2, 0.4, 0.8], np.float32), (2, 3, 3)) + expected = np.moveaxis(image * 2 - 1, -1, 0) + np.testing.assert_array_equal(small_model()._preprocess_image(image).numpy(), expected) + + def test_explicit_raw_float_counts_require_and_use_bit_depth(self): + for depth, maximum in ((8, 255), (16, 65535)): + with self.subTest(depth=depth): + image = np.broadcast_to(np.asarray([0, 128, 200], np.float32), (2, 3, 3)) + actual = small_model()._preprocess_image(image, input_range="raw", bit_depth=depth) + expected = np.moveaxis(image.astype(np.float64) / maximum * 2 - 1, -1, 0) + np.testing.assert_allclose(actual.numpy(), expected, atol=1e-7) + + def test_raw_float_without_acquisition_depth_is_rejected(self): + with self.assertRaises(ValueError): + small_model()._preprocess_image(np.full((2, 3, 3), 128.0), input_range="raw") + + def test_auto_nonunit_floats_preserve_pre_session_behavior(self): + # Before this review, only floats wholly within [0,1] were rescaled. + # Positive training-normalized values above1 passed through unchanged. + image = np.broadcast_to(np.asarray([0.2, 1.25, 1.5], np.float32), (2, 3, 3)) + np.testing.assert_array_equal(small_model()._preprocess_image(image).numpy(), + np.moveaxis(image, -1, 0)) + + def test_auto_float_range_is_decided_after_float32_rounding_as_before(self): + # float32 rounds the tiny upper overshoot to1 and the tiny negative to0. + # The old API then treats each array as unit data and applies2*x-1. + for values, expected in ( + ([0.25, 0.5, 1 + 1e-8], [-0.5, 0, 1]), + ([-1e-50, 0.25, 0.5], [-1, -0.5, 0]), + ): + with self.subTest(values=values): + image = np.asarray([[values]], dtype=np.float64) + actual = small_model()._preprocess_image(image).numpy()[:, 0, 0] + np.testing.assert_array_equal(actual, expected) + + def test_invalid_explicit_ranges_and_depths_are_rejected(self): + for options in ( + {"input_range": "other"}, {"input_range": "raw", "bit_depth": 12}, + ): + with self.subTest(options=options): + with self.assertRaises(ValueError): + small_model()._preprocess_image(np.zeros((2, 3, 3)), **options) + + def test_finite_real_inputs_are_required_including_ignored_protein_channel(self): + cases = [ + np.full((2, 3, 3), np.nan), np.full((2, 3, 3), np.inf), + np.ones((2, 3, 3), dtype=np.complex64) * (1 + 1j), + ] + ignored_target = np.zeros((2, 3, 4), np.float32) + ignored_target[..., 1] = np.nan + cases.append(ignored_target) + for image in cases: + with self.subTest(dtype=image.dtype, channels=image.shape[-1]): + with self.assertRaises(ValueError): + small_model()._preprocess_image(image) + + def test_explicit_range_limits_are_checked_before_scaling(self): + for image, options in ( + (np.full((2, 3, 3), 1.1), {"input_range": "unit"}), + (np.full((2, 3, 3), 1 + 1e-8), {"input_range": "unit"}), + (np.full((2, 3, 3), -1e-50), {"input_range": "unit"}), + (np.full((2, 3, 3), -0.1), {"input_range": "unit"}), + (np.full((2, 3, 3), 256.0), {"input_range": "raw", "bit_depth": 8}), + ): + with self.subTest(options=options): + with self.assertRaises(ValueError): + small_model()._preprocess_image(image, **options) + + def test_positive_extents_correct_channel_axis_and_numeric_dtype_are_required(self): + for image in ( + np.empty((0, 3, 3)), np.zeros((2, 3)), np.zeros((2, 3, 2)), + np.zeros((1, 2, 3, 3)), np.full((2, 3, 3), "0.5"), + np.full((2, 3, 3), 0.5, dtype=object), + ): + with self.subTest(shape=image.shape, dtype=image.dtype): + with self.assertRaises(ValueError): + small_model()._preprocess_image(image) + + +class PublicInputValidationTests(unittest.TestCase): + def test_preflight_accepts_prepared_reference_overshoot(self): + image = np.broadcast_to(np.asarray([-1, 1.0639269, 1.2511415], np.float32), (512, 512, 3)) + report = small_model().validate_inputs([image], ["A"], input_range="normalized") + self.assertTrue(report["valid"], report["errors"]) + + def test_native_512_pixel_input_passes_preflight_without_loading_weights(self): + model = small_model() + report = model.validate_inputs([np.zeros((512, 512, 3), np.uint8)], ["A"], ["C"]) + self.assertTrue(report["valid"], report["errors"]) + self.assertIsNone(model._model) + self.assertIsNone(model._vae) + + def test_non_native_geometry_is_a_preflight_error(self): + report = small_model().validate_inputs([np.zeros((256, 512, 3), np.float32)], ["A"]) + self.assertFalse(report["valid"], "Public inference requires the native 512x512 geometry") + + def test_nonfinite_input_is_a_preflight_error(self): + image = np.zeros((512, 512, 3), np.float32) + image[0, 0, 0] = np.nan + report = small_model().validate_inputs([image], ["A"]) + self.assertFalse(report["valid"]) + + def test_preflight_accepts_explicit_positive_normalized_input(self): + report = small_model().validate_inputs( + [np.full((512, 512, 3), 0.4, np.float32)], ["A"], + input_range="normalized", + ) + self.assertTrue(report["valid"], report["errors"]) + + def check_public_range(self, reliability, input_range, value, expected): + from tests.test_core_regressions import prediction_boundaries + + with prediction_boundaries() as model: + captured = [] + original = model.vae.encode + + def encode(condition): + captured.append(condition.detach().clone()) + return original(condition) + + with patch.object(model.vae, "encode", side_effect=encode): + method = model.predict_with_reliability if reliability else model.predict + options = {"num_samples": 2} if reliability else {} + result = method( + [np.full((512, 512, 3), value, np.float32)], ["A"], ["C"], + input_range=input_range, seed=23, show_progress=False, **options, + ) + self.assertEqual(len(result.images), 1) + np.testing.assert_allclose(captured[0].numpy(), expected, atol=1e-7) + + def test_predict_forwards_explicit_positive_normalized_range(self): + self.check_public_range(False, "normalized", 0.4, 0.4) + + def test_reliability_forwards_explicit_unit_range(self): + self.check_public_range(True, "unit", 0.25, -0.5) + + def test_non_native_inference_is_rejected_before_learned_components_load(self): + for reliability in (False, True): + with self.subTest(reliability=reliability): + model = small_model() + with patch.object(model, "_load_model", side_effect=AssertionError("unexpected weight load")), patch.object( + model, "_load_vae", side_effect=AssertionError("unexpected VAE load") + ): + method = model.predict_with_reliability if reliability else model.predict + with self.assertRaises(ValueError): + method([np.zeros((256, 512, 3), np.float32)], ["A"], show_progress=False) + + +class SharedTrainingRangeTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + + def load(self, image, **options): + imwrite(self.root / "cell.tif", image, photometric="minisblack") + return FinetuneDataset(str(self.root), ["cell.tif"], [1], [1], **options)[0] + + def test_dim_uint16_training_and_inference_use_the_same_dtype_scale(self): + image = np.broadcast_to(np.asarray([20, 128, 70, 200], np.uint16), (2, 3, 4)).copy() + expected = image.astype(np.float64) / 65535 * 2 - 1 + sample = self.load(image) + np.testing.assert_allclose(sample["gt_image"].numpy(), expected[..., 1:2].transpose(2, 0, 1), atol=1e-7) + inference = small_model()._preprocess_image(image).numpy() + np.testing.assert_allclose(sample["cond_image"].numpy(), inference, atol=1e-7) + + def test_training_explicit_positive_normalized_data_is_preserved(self): + image = np.broadcast_to(np.asarray([0.2, 0.3, 0.4, 0.8], np.float32), (2, 3, 4)).copy() + sample = self.load(image, input_range="normalized") + np.testing.assert_array_equal(sample["cond_image"].numpy(), image[..., [0, 2, 3]].transpose(2, 0, 1)) + np.testing.assert_array_equal(sample["gt_image"].numpy(), image[..., 1:2].transpose(2, 0, 1)) + + def test_prepared_training_tiff_preserves_reference_and_target_overshoot(self): + image = np.broadcast_to(np.asarray([-1, 1.3, 1.0639269, 1.2511415], np.float32), (2, 3, 4)).copy() + sample = self.load(image, input_range="normalized") + np.testing.assert_array_equal(sample["cond_image"].numpy(), image[..., [0, 2, 3]].transpose(2, 0, 1)) + np.testing.assert_array_equal(sample["gt_image"].numpy(), image[..., 1:2].transpose(2, 0, 1)) + + def test_training_unit_floats_and_target_are_scaled_once(self): + image = np.broadcast_to(np.asarray([0, 0.25, 0.5, 1], np.float32), (2, 3, 4)).copy() + sample = self.load(image, input_range="unit") + np.testing.assert_array_equal(sample["cond_image"].numpy()[:, 0, 0], [-1, 0, 1]) + np.testing.assert_array_equal(sample["gt_image"].numpy(), np.full((1, 2, 3), -0.5)) + + def test_training_raw_float_counts_use_explicit_sixteen_bit_depth(self): + image = np.full((2, 3, 4), 128.0, np.float32) + sample = self.load(image, input_range="raw", bit_depth=16) + np.testing.assert_allclose(sample["gt_image"].numpy(), 128 / 65535 * 2 - 1, atol=1e-7) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_resume_memory.py b/tests/test_resume_memory.py new file mode 100644 index 0000000..88013d8 --- /dev/null +++ b/tests/test_resume_memory.py @@ -0,0 +1,119 @@ +"""Bound temporary model-sized storage while restoring an EMA checkpoint. + +The training loop and checkpoint I/O are real. A large inert buffer makes the +model/EMA copies distinguishable from tiny optimizer and activation tensors, +without measuring process RSS or allocating a production-sized network. +""" + +from __future__ import annotations + +from pathlib import Path +import tempfile +import unittest +import weakref + +import numpy as np +import torch +from tifffile import imwrite +from torch.utils._python_dispatch import TorchDispatchMode +from torch.utils._pytree import tree_flatten + +from tests.test_training_controls import ControlledDenoiser, TrainingHarness + + +class _StorageBudget(TorchDispatchMode): + """Observe distinct live large storages without owning their tensors.""" + + def __init__(self, minimum_bytes): + super().__init__() + self.minimum_bytes = minimum_bytes + self.references = {} + self.peak_bytes = 0 + self.frozen = False + + def observe(self, value): + if self.frozen: + return + for tensor in tree_flatten(value)[0]: + if not isinstance(tensor, torch.Tensor) or tensor.device.type != "cpu": + continue + storage = tensor.untyped_storage() + if storage.nbytes() >= self.minimum_bytes: + key = (storage.data_ptr(), storage.nbytes()) + self.references.setdefault(key, []).append(weakref.ref(tensor)) + live_bytes = 0 + for key, references in list(self.references.items()): + live = [] + for reference in references: + tensor = reference() + if tensor is not None: + storage = tensor.untyped_storage() + if (storage.data_ptr(), storage.nbytes()) == key: + live.append(reference) + if live: + self.references[key] = live + live_bytes += key[1] + else: + del self.references[key] + self.peak_bytes = max(self.peak_bytes, live_bytes) + + def __torch_dispatch__(self, function, types, args=(), kwargs=None): + result = function(*args, **(kwargs or {})) + self.observe(result) + return result + + def freeze(self): + self.frozen = True + + +class _BufferedDenoiser(ControlledDenoiser): + def __init__(self): + super().__init__() + self.register_buffer("memory_probe", torch.full((65536,), 0.75, dtype=torch.float32)) + + +class ResumeMemoryTests(unittest.TestCase): + def test_ema_resume_has_no_discarded_full_model_shadow(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + images = root / "images" + images.mkdir() + filename = "cell.tif" + imwrite(images / filename, np.full((2, 2, 4), 0.2, np.float32), + photometric="minisblack") + + def harness(output): + run = TrainingHarness(images, [filename], root / output) + run.network = _BufferedDenoiser().eval().requires_grad_(False) + run.model._model = run.network + return run + + checkpoint = harness("checkpoint").run(use_ema=True) + resumed = harness("resumed") + buffer_bytes = resumed.network.memory_probe.nbytes + observer = _StorageBudget(buffer_bytes) + observer.observe(resumed.network.memory_probe) + # The budget applies to restoring state, before activations and + # subsequent training/checkpoint exports introduce other lifetimes. + hook = resumed.network.register_forward_pre_hook( + lambda _module, _inputs: observer.freeze() + ) + try: + with observer: + resumed.run(use_ema=True, num_epochs=2, + resume_from=str(checkpoint.output)) + finally: + hook.remove() + + self.assertTrue(observer.frozen, "The resumed training did not reach its first forward pass") + # CPU resume can retain the live network, raw checkpoint network, + # and its saved EMA. These three distinct buffers are sufficient; + # an additional full shadow, even transiently, is unnecessary. + self.assertLessEqual( + observer.peak_bytes, 3 * buffer_bytes, + "EMA resume exceeded the live model + raw checkpoint + saved EMA storage budget", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_training_controls.py b/tests/test_training_controls.py new file mode 100644 index 0000000..c5ae288 --- /dev/null +++ b/tests/test_training_controls.py @@ -0,0 +1,321 @@ +"""Tiny CPU integrations for accumulation, resumable state, EMA, and AMP. + +The optimizer, autograd, TIFF dataset, checkpoint files, and training loop are +real. Only learned networks, the fixed sigma schedule, and a deterministic LR +factory are controlled. Accumulation gradients have an independent algebraic +reference; resume is compared to an uninterrupted stochastic training run. +""" + +from __future__ import annotations + +import contextlib +import functools +import importlib +import io +import json +import math +from pathlib import Path +import random +import tempfile +import types +import unittest +from unittest.mock import patch +import warnings + +import numpy as np +import torch +from tifffile import imwrite + +from proticelli import Model, _training + + +class ControlledDenoiser(torch.nn.Module): + def __init__(self, *, analytic=False, fail_on=None): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([[0.1]], dtype=torch.float32)) + self.register_buffer("forward_count", torch.tensor(0, dtype=torch.int64)) + self.analytic = analytic + self.fail_on = fail_on + self.calls = 0 + self.traces = [] + self.projection_dtypes = [] + self.parameter_dtypes = [] + + def forward(self, image, timestep, **_kwargs): + self.calls += 1 + if self.calls == self.fail_on: + raise InterruptedError("controlled interruption at a model boundary") + self.forward_count.add_(1) + self.traces.append({ + "target": image[:, 16:].detach().float().mean(dim=(1, 2, 3)).clone(), + "input": image[:, :16].detach().float().clone(), + }) + ones = torch.ones_like(image[:, :1]).movedim(1, -1) + projection = torch.nn.functional.linear(ones, self.weight).movedim(-1, 1) + self.projection_dtypes.append(projection.dtype) + self.parameter_dtypes.append(self.weight.dtype) + prediction = projection.repeat(1, 16, 1, 1) + if self.analytic: + # At sigma=.5, c_in=sqrt(2), c_skip=.5, c_out=1/(2*sqrt(2)). + # F=w-c_in*x cancels the noisy contribution in D=c_skip*x+c_out*F. + prediction = prediction - image[:, :16] + else: + prediction = prediction + 0.01 * timestep[:, None, None, None] + return types.SimpleNamespace(sample=prediction) + + def save_pretrained(self, directory, **_kwargs): + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + torch.save(self.state_dict(), directory / "test_weights.pt") + (directory / "config.json").write_text("{}", encoding="utf-8") + + +class ControlledVAE(torch.nn.Module): + config = types.SimpleNamespace(scaling_factor=4.0) + + def __init__(self): + super().__init__() + self.register_buffer("anchor", torch.tensor(0.0)) + + @property + def dtype(self): + return self.anchor.dtype + + def encode(self, image): + mean = image.mean(dim=1, keepdim=True).repeat(1, 16, 1, 1) + return types.SimpleNamespace(latent_dist=types.SimpleNamespace(sample=lambda: mean.clone())) + + +class QuietProgress: + def __init__(self, iterable): + self.iterable = iterable + + def __iter__(self): + return iter(self.iterable) + + def set_postfix(self, **_kwargs): + pass + + +class TrainingHarness: + def __init__(self, images, names, output, *, analytic=False, fail_on=None, dtype="float32"): + self.images, self.names, self.output = images, names, output + self.model = Model(device="cpu", dtype=dtype, protein_map={"P": 1}, cellline_map={"C": 1}) + self.network = ControlledDenoiser(analytic=analytic, fail_on=fail_on) + self.network.to(self.model.dtype).eval().requires_grad_(False) + self.model._model = self.network + self.model._vae = ControlledVAE().to(self.model.dtype).eval().requires_grad_(False) + self.optimizers, self.schedules, self.steps = [], [], [] + + def run(self, **overrides): + options = { + "num_epochs": 1, "batch_size": 2, "gradient_accumulation_steps": 1, + "learning_rate": 0.01, "lr_scheduler_type": "linear", "lr_warmup_steps": 0, + "checkpointing_steps": 1, "save_model_epochs": 1, "num_workers": 0, + "max_grad_norm": 1e6, "adam_weight_decay": 0, + "input_range": "normalized", "seed": 139, + } + options.update(overrides) + optimizer_class = torch.optim.AdamW + + def optimizer_factory(*args, **kwargs): + optimizer = optimizer_class(*args, **kwargs) + self.optimizers.append(optimizer) + return optimizer + + def lr_factory(_name, *, optimizer, **_kwargs): + # A real stateful schedule with nonconstant LR makes missing LR + # restoration observable when comparing resumed optimization. + schedule = torch.optim.lr_scheduler.LambdaLR( + optimizer, lr_lambda=lambda step: max(0.1, 1 - step / 20) + ) + original_step = optimizer.step + + @functools.wraps(original_step) + def step(*args, **kwargs): + record = { + "before": self.network.weight.detach().float().clone(), + "gradient": self.network.weight.grad.detach().float().clone(), + "through_trace": len(self.network.traces), + } + result = original_step(*args, **kwargs) + record["after"] = self.network.weight.detach().float().clone() + self.steps.append(record) + return result + + optimizer.step = step + self.schedules.append(schedule) + return schedule + + optimization = importlib.import_module("diffusers.optimization") + edm = importlib.import_module("proticelli.schedulers.edm_scheduler") + numpy_state, python_state = np.random.get_state(), random.getstate() + try: + with contextlib.ExitStack() as stack: + stack.enter_context(patch.object(_training.torch.optim, "AdamW", optimizer_factory)) + stack.enter_context(patch.object(optimization, "get_scheduler", lr_factory)) + stack.enter_context(patch.object(edm, "create_edm_scheduler", lambda **_kwargs: types.SimpleNamespace( + config=types.SimpleNamespace(num_train_timesteps=1000), + sigmas=torch.full((1001,), 0.5), + ))) + stack.enter_context(patch.object(_training, "tqdm", lambda iterable, **_kwargs: QuietProgress(iterable))) + stack.enter_context(torch.random.fork_rng(devices=[])) + stack.enter_context(contextlib.redirect_stdout(io.StringIO())) + stack.enter_context(warnings.catch_warnings()) + warnings.filterwarnings("ignore", category=FutureWarning, message=".*GradScaler.*") + # Intentionally unrelated initial state. Explicit seed or resumed + # checkpoint RNG must control the actual training draws. + torch.random.default_generator.manual_seed(991) + self.model.fit( + self.images, self.names, ["P"] * len(self.names), ["C"] * len(self.names), + output_dir=str(self.output), **options, + ) + finally: + np.random.set_state(numpy_state) + random.setstate(python_state) + return self + + def state(self, path=None): + path = self.output / "training_state.pt" if path is None else Path(path) + return torch.load(path, map_location="cpu", weights_only=False) + + +class TrainingControlTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.images = self.root / "images" + self.images.mkdir() + self.names = [] + for index, value in enumerate((-0.6, -0.1, 0.2, 0.7, 0.9, 0.4)): + filename = f"cell_{index:03d}.tif" + imwrite(self.images / filename, np.full((2, 2, 4), value, np.float32), photometric="minisblack") + self.names.append(filename) + + def harness(self, name, n=6, **kwargs): + return TrainingHarness(self.images, self.names[:n], self.root / name, **kwargs) + + def assert_gradient_reference(self, run): + previous = 0 + for step in run.steps: + groups = run.network.traces[previous:step["through_trace"]] + targets = torch.cat([group["target"] for group in groups]).double().numpy() + # With the analytic denoiser, D=w/(2*sqrt(2)), weight(sigma)=8. + # The sample-mean squared loss therefore has dL/dw=2w-4sqrt(2)*mean(y). + expected = 2 * step["before"].double().numpy() - 4 * math.sqrt(2) * targets.mean() + np.testing.assert_allclose(step["gradient"].numpy(), expected, rtol=2e-6, atol=2e-6) + previous = step["through_trace"] + + def test_accumulation_weights_unequal_minibatches_by_samples(self): + run = self.harness("unequal", n=3, analytic=True).run( + batch_size=2, gradient_accumulation_steps=2 + ) + self.assertEqual(len(run.steps), 1) + self.assert_gradient_reference(run) + + def test_accumulation_tail_uses_its_actual_sample_count(self): + run = self.harness("tail", n=5, analytic=True).run( + batch_size=2, gradient_accumulation_steps=2 + ) + self.assertEqual(len(run.steps), 2) + self.assert_gradient_reference(run) + + def test_seeded_fresh_training_has_a_reproducible_passing_control(self): + first = self.harness("first", n=4).run() + second = self.harness("second", n=4).run() + torch.testing.assert_close(first.state()["model"]["weight"], second.state()["model"]["weight"], rtol=0, atol=0) + for a, b in zip(first.network.traces, second.network.traces): + torch.testing.assert_close(a["input"], b["input"], rtol=0, atol=0) + + def test_mid_epoch_resume_restores_model_optimizer_lr_cursor_and_rng(self): + uninterrupted = self.harness("uninterrupted").run(num_epochs=2, use_ema=True, ema_decay=0.5) + interrupted = self.harness("interrupted", fail_on=3) + with self.assertRaises(InterruptedError): + interrupted.run(num_epochs=2, use_ema=True, ema_decay=0.5) + checkpoint = interrupted.output / "checkpoint-2" + partial = interrupted.state(checkpoint / "training_state.pt") + self.assertEqual((partial["epoch"], partial["next_batch"], partial["global_step"]), (0, 2, 2)) + resumed = self.harness("resumed").run( + num_epochs=2, resume_from=str(checkpoint), use_ema=True, ema_decay=0.5 + ) + expected, actual = uninterrupted.state(), resumed.state() + required = {"version", "model", "optimizer", "lr_scheduler", "scaler", "ema", + "epoch", "next_batch", "global_step", "data_seed", "rng", "configuration"} + self.assertTrue(required.issubset(actual)) + self.assertEqual((actual["epoch"], actual["next_batch"], actual["global_step"]), (2, 0, 6)) + for name in expected["model"]: + torch.testing.assert_close(actual["model"][name], expected["model"][name], rtol=0, atol=0) + torch.testing.assert_close(actual["ema"][name], expected["ema"][name], rtol=0, atol=0) + self.assertEqual(actual["lr_scheduler"]["last_epoch"], expected["lr_scheduler"]["last_epoch"]) + self.assertEqual(actual["optimizer"]["param_groups"], expected["optimizer"]["param_groups"]) + for parameter, state in expected["optimizer"]["state"].items(): + for key, value in state.items(): + torch.testing.assert_close(actual["optimizer"]["state"][parameter][key], value, rtol=0, atol=0) + torch.testing.assert_close(actual["rng"]["torch_cpu"], expected["rng"]["torch_cpu"], rtol=0, atol=0) + self.assertEqual(len(resumed.network.traces), len(uninterrupted.network.traces) - 2) + for a, b in zip(resumed.network.traces, uninterrupted.network.traces[2:]): + torch.testing.assert_close(a["input"], b["input"], rtol=0, atol=0) + + def test_resume_accepts_a_training_state_file_path(self): + complete = self.harness("complete", n=4).run() + checkpoint = complete.output / "checkpoint-1" / "training_state.pt" + resumed = self.harness("file_resume", n=4).run(resume_from=str(checkpoint)) + torch.testing.assert_close(resumed.state()["model"]["weight"], complete.state()["model"]["weight"], rtol=0, atol=0) + + def test_ema_updates_follow_the_declared_recurrence_and_export_separately(self): + decay = 0.5 + run = self.harness("ema", n=4).run(use_ema=True, ema_decay=decay) + expected = torch.tensor([[0.1]], dtype=torch.float32) + for step in run.steps: + expected = decay * expected + (1 - decay) * step["after"] + state = run.state() + torch.testing.assert_close(state["ema"]["weight"], expected) + torch.testing.assert_close(run.model.model.weight, expected) + raw = torch.load(run.output / "unet" / "test_weights.pt", weights_only=True) + ema = torch.load(run.output / "unet_ema" / "test_weights.pt", weights_only=True) + torch.testing.assert_close(raw["weight"], state["model"]["weight"]) + torch.testing.assert_close(ema["weight"], expected) + self.assertEqual(state["ema"]["forward_count"].dtype, torch.int64) + self.assertEqual(state["ema"]["forward_count"].item(), len(run.network.traces)) + torch.testing.assert_close(ema["forward_count"], raw["forward_count"]) + self.assertFalse(torch.equal(raw["weight"], ema["weight"])) + metadata = json.loads((run.output / "proticelli_checkpoint.json").read_text()) + self.assertEqual(metadata["inference_subfolder"], "unet_ema") + + def test_bf16_cpu_uses_autocast_with_fp32_master_parameters(self): + run = self.harness("bf16", n=4).run(mixed_precision="bf16") + self.assertTrue(all(dtype == torch.bfloat16 for dtype in run.network.projection_dtypes)) + self.assertTrue(all(dtype == torch.float32 for dtype in run.network.parameter_dtypes)) + self.assertEqual(run.state()["model"]["weight"].dtype, torch.float32) + self.assertTrue(torch.isfinite(run.network.weight).all().item()) + + def test_fp16_cpu_is_rejected_before_a_training_forward(self): + run = self.harness("invalid_fp16", n=4) + with self.assertRaisesRegex(ValueError, "(?i)fp16|cuda"): + run.run(mixed_precision="fp16") + self.assertEqual(run.network.calls, 0) + self.assertFalse(run.network.training) + self.assertFalse(run.network.weight.requires_grad) + + def test_half_precision_inference_model_trains_in_fp32_and_restores_dtype(self): + run = self.harness("half_inference", n=4, dtype="float16").run() + self.assertTrue(all(dtype == torch.float32 for dtype in run.network.parameter_dtypes)) + self.assertEqual(run.state()["model"]["weight"].dtype, torch.float32) + self.assertEqual(run.network.weight.dtype, torch.float16) + self.assertEqual(run.model.vae.dtype, torch.float16) + self.assertFalse(run.network.training) + self.assertFalse(run.network.weight.requires_grad) + + def test_training_failure_restores_inference_dtype_eval_and_frozen_parameters(self): + run = self.harness("failed_half", n=4, dtype="float16", fail_on=1) + with self.assertRaises(InterruptedError): + run.run() + self.assertEqual(run.network.weight.dtype, torch.float16) + self.assertEqual(run.model.vae.dtype, torch.float16) + self.assertFalse(run.network.training) + self.assertFalse(run.network.weight.requires_grad) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_training_regressions.py b/tests/test_training_regressions.py new file mode 100644 index 0000000..9b85e64 --- /dev/null +++ b/tests/test_training_regressions.py @@ -0,0 +1,255 @@ +"""Test-first checks for core fine-tuning behavior; production is unchanged. + +Numerical references come from TIFF dtype ranges and the algebraic EDM loss +weight. Integration cases exercise the actual data loader, autograd, optimizer, +and training loop with tiny replacements for the learned VAE and denoiser. +Known defects deliberately fail; they are not marked as expected failures. +""" + +from __future__ import annotations + +import contextlib +import importlib +import io +import math +from pathlib import Path +import tempfile +import types +import unittest +from unittest.mock import Mock, patch +import warnings + +import numpy as np +from numpy.testing import assert_allclose +import torch +from tifffile import imwrite + +from proticelli import _training +from proticelli.utils.edm_utils import edm_loss_weight + + +class FineTuneInputTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + + def load(self, array): + name = 'cell.tif' + imwrite(self.root / name, array, photometric='minisblack') + return _training.FinetuneDataset(str(self.root), [name], [7], [3])[0] + + def assert_normalization(self, dtype, pixels): + raw = np.broadcast_to(np.asarray(pixels, dtype=dtype), (3, 5, 4)).copy() + sample = self.load(raw) + # A raw integer count is normalized using its acquisition dtype, not + # the brightest observed pixel. Channel1 is the target, 0/2/3 condition. + expected = raw.astype(np.float64) / np.iinfo(dtype).max * 2 - 1 + assert_allclose(sample['cond_image'].numpy(), expected[..., [0, 2, 3]].transpose(2, 0, 1), atol=1e-7) + assert_allclose(sample['gt_image'].numpy(), expected[..., 1:2].transpose(2, 0, 1), atol=1e-7) + self.assertEqual(sample['label'], 7) + self.assertEqual(sample['cell_line'], 3) + + def test_raw_uint8_with_nontrivial_channel_values(self): + self.assert_normalization(np.uint8, [20, 255, 80, 160]) + + def test_raw_uint16_with_full_scale_reference(self): + self.assert_normalization(np.uint16, [200, 65535, 2048, 40000]) + + def test_dim_uint16_retains_its_sixteen_bit_scale(self): + self.assert_normalization(np.uint16, [20, 128, 70, 200]) + + def test_one_count_uint8_is_not_a_unit_range_float(self): + self.assert_normalization(np.uint8, [0, 1, 1, 0]) + + def test_unit_range_float_tiff_is_scaled_once(self): + raw = np.broadcast_to(np.array([0.0, 0.25, 0.5, 1.0], np.float32), (3, 5, 4)).copy() + sample = self.load(raw) + assert_allclose(sample['cond_image'].numpy()[:, 0, 0], [-1, 0, 1], atol=1e-7) + assert_allclose(sample['gt_image'].numpy(), -0.5, atol=1e-7) + + def test_signed_float_tiff_keeps_reference_order_and_target(self): + raw = np.broadcast_to(np.array([-0.8, 0.25, 0.4, 0.6], np.float32), (3, 5, 4)).copy() + sample = self.load(raw) + assert_allclose(sample['cond_image'].numpy()[:, 0, 0], [-0.8, 0.4, 0.6], atol=1e-7) + assert_allclose(sample['gt_image'].numpy(), 0.25, atol=1e-7) + + +class EDMLossWeightTests(unittest.TestCase): + def check_weights(self, dtype): + sigma = torch.tensor([0.002, 0.5, 80.0], dtype=dtype) + # (s^2+d^2)/(s*d)^2 = 1/s^2 + 1/d^2. Use represented sigma + # values in a float64 reference, so half quantization is accounted for. + expected = 1 / sigma.double().numpy() ** 2 + 4 + actual = edm_loss_weight(sigma) + self.assertTrue(torch.isfinite(actual).all().item(), f'Non-finite weights: {actual}') + assert_allclose(actual.double().numpy(), expected, rtol=5e-7, atol=1e-7) + + def test_float32_weights_match_independent_algebraic_reference(self): + self.check_weights(torch.float32) + + def test_half_precision_input_still_yields_finite_accurate_loss_weights(self): + self.check_weights(torch.float16) + + +class _TinyDenoiser(torch.nn.Module): + """Consumes one timestep per sample, as the real timestep embedding does.""" + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(0.1)) + self.forward_calls = 0 + self.saved_weights = [] + + def forward(self, image, timestep, **kwargs): + self.forward_calls += 1 + # Normal batched embedding/broadcasting requires a one-dimensional + # timestep vector. There is no assertion tailored to the implementation. + temporal = timestep[:, None, None, None] + prediction = self.weight * torch.ones_like(image[:, :16]) + temporal * 0.01 + return types.SimpleNamespace(sample=prediction) + + def save_pretrained(self, path): + self.saved_weights.append((str(path), self.weight.detach().clone())) + + +class _TinyVAE(torch.nn.Module): + config = types.SimpleNamespace(scaling_factor=4.0) + + def encode(self, image): + latent = image.mean(dim=1, keepdim=True).repeat(1, 16, 1, 1) + return types.SimpleNamespace(latent_dist=types.SimpleNamespace(sample=lambda: latent.clone())) + + +class _QuietProgress: + def __init__(self, iterable): + self.iterable = iterable + + def __iter__(self): + return iter(self.iterable) + + def set_postfix(self, **kwargs): + pass + + +class _StatefulSchedule: + """Constant-rate scheduler boundary that supports actual checkpoint state.""" + + def __init__(self, optimizer): + self.optimizer = optimizer + self.steps = 0 + self.step = Mock(wraps=self._step) + + def _step(self): + self.steps += 1 + + def get_last_lr(self): + return [group['lr'] for group in self.optimizer.param_groups] + + def state_dict(self): + return {'steps': self.steps} + + def load_state_dict(self, state): + self.steps = state['steps'] + + +class FineTuneLoopTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.images = self.root / 'images' + self.images.mkdir() + self.optimizers = [] + self.schedules = [] + + def run_loop(self, n_images, *, batch_size, accumulation=1, **options): + names = [] + for index in range(n_images): + name = f'cell_{index}.tif' + image = np.broadcast_to(np.array([-0.6, -0.2, 0.3, 0.7], np.float32), (2, 2, 4)).copy() + imwrite(self.images / name, image, photometric='minisblack') + names.append(name) + network = _TinyDenoiser() + model = types.SimpleNamespace( + device=torch.device('cpu'), dtype=torch.float32, model=network, vae=_TinyVAE(), + protein_map={'P': 1}, cellline_map={'C': 1}, _resolve_protein_name=lambda name: name, + ) + original_optimizer = torch.optim.AdamW + + def optimizer_factory(*args, **kwargs): + optimizer = original_optimizer(*args, **kwargs) + optimizer.step = Mock(wraps=optimizer.step) + self.optimizers.append(optimizer) + return optimizer + + def schedule_factory(name, *, optimizer, **kwargs): + schedule = _StatefulSchedule(optimizer) + self.schedules.append(schedule) + return schedule + + optimization_module = importlib.import_module('diffusers.optimization') + scheduler_module = importlib.import_module('proticelli.schedulers.edm_scheduler') + + def noise_schedule_factory(**kwargs): + return types.SimpleNamespace( + config=types.SimpleNamespace(num_train_timesteps=1000), + sigmas=torch.full((1001,), 0.5), + ) + + with contextlib.ExitStack() as stack: + stack.enter_context(patch.object(optimization_module, 'get_scheduler', schedule_factory)) + stack.enter_context(patch.object(scheduler_module, 'create_edm_scheduler', noise_schedule_factory)) + stack.enter_context(patch.object(_training.torch.optim, 'AdamW', optimizer_factory)) + stack.enter_context(patch.object(_training, 'tqdm', lambda iterable, **kwargs: _QuietProgress(iterable))) + stack.enter_context(torch.random.fork_rng(devices=[])) + stack.enter_context(contextlib.redirect_stdout(io.StringIO())) + stack.enter_context(warnings.catch_warnings()) + warnings.filterwarnings('ignore', message=".*pin_memory.*", category=UserWarning) + torch.random.default_generator.manual_seed(19) + _training.run_finetuning( + model, str(self.images), names, ['P'] * n_images, ['C'] * n_images, + num_epochs=1, batch_size=batch_size, gradient_accumulation_steps=accumulation, + num_workers=0, lr_warmup_steps=0, learning_rate=1e-3, + output_dir=str(self.root / 'output'), checkpointing_steps=1000, + **options, + ) + return model, network + + def test_regular_batches_update_parameters_and_finish_in_eval_mode(self): + model, network = self.run_loop(4, batch_size=2) + self.assertEqual(network.forward_calls, 2) + self.assertEqual(self.optimizers[0].step.call_count, 2) + self.assertEqual(self.schedules[0].step.call_count, 2) + self.assertTrue(math.isfinite(network.weight.item())) + self.assertNotEqual(network.weight.item(), torch.tensor(0.1).item()) + self.assertFalse(network.training) + self.assertFalse(network.weight.requires_grad) + self.assertEqual(model.checkpoint_dir, self.root / 'output') + self.assertTrue(network.saved_weights) + + def test_batch_size_one_is_a_valid_training_batch(self): + _, network = self.run_loop(1, batch_size=1) + self.assertEqual(network.forward_calls, 1) + self.assertEqual(self.optimizers[0].step.call_count, 1) + self.assertTrue(torch.isfinite(network.weight).item()) + + def test_final_singleton_remainder_is_trained_without_shape_failure(self): + _, network = self.run_loop(3, batch_size=2) + self.assertEqual(network.forward_calls, 2) + self.assertEqual(self.optimizers[0].step.call_count, 2) + + def test_gradient_accumulation_updates_once_per_group_including_remainder(self): + self.run_loop(10, batch_size=2, accumulation=2) + # Ten samples / two samples per minibatch = five minibatches. + # Two minibatches per update require three updates, including the tail. + self.assertEqual(self.optimizers[0].step.call_count, 3) + self.assertEqual(self.schedules[0].step.call_count, 3) + + def test_missing_resume_checkpoint_cannot_silently_start_training(self): + with self.assertRaises((FileNotFoundError, NotImplementedError)): + self.run_loop(2, batch_size=2, resume_from=str(self.root / 'missing-checkpoint')) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_transformer_compatibility.py b/tests/test_transformer_compatibility.py new file mode 100644 index 0000000..70a964e --- /dev/null +++ b/tests/test_transformer_compatibility.py @@ -0,0 +1,143 @@ +"""Compare inherited normalization behavior with its upstream implementation.""" + +import unittest +import tempfile +import json +from pathlib import Path + +import torch +from diffusers.models.attention import BasicTransformerBlock as UpstreamBlock + +from proticelli.models.basic_transformer_block import BasicTransformerBlock +from proticelli.models.dit import DiTTransformer2DModel +from proticelli import Model + + +class TransformerCompatibilityTests(unittest.TestCase): + def test_legacy_discrete_loading_preserves_weights_and_reserved_null_row(self): + with torch.random.fork_rng(devices=[]), tempfile.TemporaryDirectory() as directory: + torch.random.default_generator.manual_seed(84) + legacy = DiTTransformer2DModel( + num_attention_heads=2, attention_head_dim=4, in_channels=4, + out_channels=2, num_layers=1, sample_size=8, patch_size=2, + num_embeds_ada_norm=16, num_protein_labels=20, num_cell_labels=3, + norm_type="ada_norm_zero", legacy_discrete_capacity=True, + ).eval() + path = Path(directory) / "unet" + legacy.save_pretrained(path) + config_path = path / "config.json" + config = json.loads(config_path.read_text()) + # An old checkpoint predates this option; its actual weights and + # capacity remain unchanged, including the reserved dropout row. + config.pop("legacy_discrete_capacity") + config_path.write_text(json.dumps(config), encoding="utf-8") + explicit = DiTTransformer2DModel.from_pretrained( + path, local_files_only=True, legacy_discrete_capacity=True).eval() + automatic = Model(checkpoint_dir=directory, device="cpu", dtype="float32", + protein_map={"KNOWN": 1}, cellline_map={"CELL": 1}).model + for name, value in legacy.state_dict().items(): + torch.testing.assert_close(explicit.state_dict()[name], value, rtol=0, atol=0) + torch.testing.assert_close(automatic.state_dict()[name], value, rtol=0, atol=0) + image = torch.zeros(2, 4, 8, 8) + conditions = {"timestep": torch.tensor([2, 5]), + "protein_labels": torch.tensor([0, 15]), + "cell_line_labels": torch.tensor([0, 2])} + with torch.no_grad(): + expected = legacy(image, **conditions).sample + torch.testing.assert_close(automatic(image, **conditions).sample, expected, rtol=0, atol=0) + conditions["protein_labels"] = torch.tensor([0, 19]) + with self.assertRaisesRegex(ValueError, "no trained class embeddings"): + automatic(image, **conditions) + + def test_discrete_dit_accepts_the_full_declared_protein_vocabulary(self): + with torch.random.fork_rng(devices=[]), tempfile.TemporaryDirectory() as directory: + torch.random.default_generator.manual_seed(83) + model = DiTTransformer2DModel( + num_attention_heads=2, attention_head_dim=4, in_channels=4, + out_channels=2, num_layers=1, sample_size=8, patch_size=2, + num_embeds_ada_norm=16, num_protein_labels=20, num_cell_labels=3, + norm_type="ada_norm_zero", + ).eval() + image = torch.zeros(2, 4, 8, 8) + conditions = {"timestep": torch.tensor([2, 5]), + "protein_labels": torch.tensor([0, model.config.num_protein_labels - 1]), + "cell_line_labels": torch.tensor([0, model.config.num_cell_labels - 1])} + with torch.no_grad(): + expected = model(image, **conditions).sample + self.assertEqual(expected.shape, (2, 2, 8, 8)) + self.assertTrue(torch.isfinite(expected).all().item()) + model.save_pretrained(directory) + restored = DiTTransformer2DModel.from_pretrained(directory, local_files_only=True).eval() + torch.testing.assert_close(restored(image, **conditions).sample, expected, rtol=0, atol=0) + + def test_discrete_and_single_norms_match_upstream_for_the_same_weights(self): + with torch.random.fork_rng(devices=[]): + torch.random.default_generator.manual_seed(81) + hidden = torch.linspace(-1, 1, 2 * 4 * 8).reshape(2, 4, 8) + for mode in ("ada_norm_zero", "ada_norm_single"): + with self.subTest(mode=mode): + options = dict(dim=8, num_attention_heads=2, attention_head_dim=4, + num_embeds_ada_norm=16, norm_type=mode) + reference = UpstreamBlock(**options).eval() + actual = BasicTransformerBlock(**options).eval() + actual.load_state_dict(reference.state_dict()) + conditions = {"timestep": torch.tensor([2, 5]), + "class_labels": torch.tensor([1, 2])} + if mode == "ada_norm_single": + conditions["timestep"] = torch.linspace(-0.3, 0.7, 2 * 6 * 8).reshape(2, 48) + with torch.no_grad(): + expected = reference(hidden, **conditions) + observed = actual(hidden, **conditions) + torch.testing.assert_close(observed, expected, rtol=0, atol=0) + + def test_dit_forwards_discrete_timestep_and_protein_labels_to_blocks(self): + with torch.random.fork_rng(devices=[]): + torch.random.default_generator.manual_seed(82) + model = DiTTransformer2DModel( + num_attention_heads=2, attention_head_dim=4, in_channels=4, + out_channels=2, num_layers=1, sample_size=8, patch_size=2, + num_embeds_ada_norm=16, num_protein_labels=5, num_cell_labels=3, + norm_type="ada_norm_zero", + ).eval() + image = torch.linspace(-1, 1, 2 * 4 * 8 * 8).reshape(2, 4, 8, 8) + conditions = {"timestep": torch.tensor([2, 5]), + "protein_labels": torch.tensor([1, 2]), + "cell_line_labels": torch.tensor([1, 2])} + # The actual upstream block supplies a separate executable reference + # for conditioning; model embedding/output geometry stays real. + options = dict(dim=8, num_attention_heads=2, attention_head_dim=4, + num_embeds_ada_norm=16, norm_type="ada_norm_zero", + norm_elementwise_affine=model.config.norm_elementwise_affine, + norm_eps=model.config.norm_eps, + activation_fn=model.config.activation_fn, + attention_bias=model.config.attention_bias) + reference = UpstreamBlock(**options).eval() + reference.load_state_dict(model.transformer_blocks[0].state_dict()) + with torch.no_grad(): + embedding = torch.cat([ + model.emb_protein_labels(conditions["timestep"], conditions["protein_labels"]), + model.emb_cell_line_labels(conditions["timestep"], conditions["cell_line_labels"]), + ], dim=1) + patches = model.pos_embed(image) + expected = reference(patches, timestep=conditions["timestep"], + class_labels=conditions["protein_labels"]) + observed = model.transformer_blocks[0](patches, embedding=embedding, + timestep=conditions["timestep"], class_labels=conditions["protein_labels"]) + torch.testing.assert_close(observed, expected, rtol=0, atol=0) + output = model(image, **conditions).sample + self.assertEqual(output.shape, (2, 2, 8, 8)) + self.assertTrue(torch.isfinite(output).all().item()) + original_input = image.clone().requires_grad_(True) + original = model(original_input, **conditions).sample + original_gradient, = torch.autograd.grad(original.square().mean(), original_input) + model.gradient_checkpointing = True + checkpointed_input = image.clone().requires_grad_(True) + checkpointed = model(checkpointed_input, **conditions).sample + checkpointed_gradient, = torch.autograd.grad( + checkpointed.square().mean(), checkpointed_input) + torch.testing.assert_close(checkpointed, original, rtol=0, atol=0) + torch.testing.assert_close(checkpointed_gradient, original_gradient, rtol=0, atol=0) + + +if __name__ == "__main__": + unittest.main()