| --- |
| license: apache-2.0 |
| library_name: stip |
| tags: |
| - jax |
| - flax |
| - stochastic-interpolants |
| - tutorial |
| --- |
| |
| # STIP tutorial checkpoints |
|
|
| Small checkpoints used by the [`stip`](https://github.com/instadeepai/stip) tutorial |
| notebooks, so that a tutorial can demonstrate sampling without spending ten |
| minutes training first. They are toy models (a two-layer MLP, 18k-25k parameters, |
| trained for 3000 steps on a 4-component 2D Gaussian mixture) and have no |
| use outside the notebooks. |
|
|
| Checkpoints are [Orbax](https://orbax.readthedocs.io) directories written by |
| `stip`'s own `TrainingIOHandler`, holding `params`, `opt_state`, `ema_params` and |
| `extra` (EMA decay and step count) as separately-restorable items. |
|
|
| ## `conditioning_and_guidance/` |
|
|
| Used by `tutorials/notebooks/4.conditioning_and_guidance.ipynb`. Both models are |
| `VelocityOneSidedGenerativeModel`s with a `FlowMatchingOneSidedInterpolant`, but over |
| different modalities: |
|
|
| | Path | Model | Modalities | Role in the notebook | |
| |---|---|---|---| |
| | `conditioning_and_guidance/joint_model` | Unconditional cross-modal MLP | `coordinates` (continuous, 2D) and `index` (discrete, 4 categories) | Intrinsic guidance (Section 3): conditioning a model that was never trained to be conditional | |
| | `conditioning_and_guidance/context_model` | The same MLP plus a label context path, trained with 50% context dropout | `coordinates` only; the corner label is passed as `context_data` instead of as a modality | Context conditioning and classifier-free guidance (Sections 4-5) | |
|
|
| ### Loading |
|
|
| ```python |
| from flax import nnx |
| from huggingface_hub import snapshot_download |
| from stip.training.checkpointer import Checkpointer, CheckpointerConfig |
| |
| path = snapshot_download( |
| "InstaDeepAI/STIP-tutorials", allow_patterns="conditioning_and_guidance/joint_model/*" |
| ) |
| gen_model = ... # build the same model structure as the notebook |
| graphdef, params = nnx.split(gen_model, nnx.Param) |
| checkpointer = Checkpointer( |
| CheckpointerConfig( |
| checkpoint_dir=f"{path}/conditioning_and_guidance/joint_model", |
| max_to_keep=None, # read-only: never mutate a downloaded directory |
| ) |
| ) |
| gen_model = nnx.merge(graphdef, checkpointer.restore_ema(params)) |
| ``` |
|
|
| `restore_ema` reads only `ema_params` and `extra`, and applies the same bias |
| correction the training loop uses for evaluation. |
|
|
| ## Reproducing |
|
|
| ```bash |
| uv run python tutorials/scripts/train_conditioning_checkpoints.py |
| ``` |
|
|
| The script mirrors the notebook's model definitions and PRNG chain, so it |
| reproduces these exact weights. A checkpoint pins the parameter structure: if a |
| notebook's network changes, re-run the script and re-upload. |
|
|
