Add minimal build CI and runtime regression fixes

.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

README.md

@@ -118,7 +118,8 @@ python -m proticelli_web --mode real
 ```
 
 Open <http://127.0.0.1:8000>. 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:

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"])

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":

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:

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
         -------

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,13 +98,16 @@ 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__()
 
         self.attention_head_dim = attention_head_dim
         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

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: