File size: 2,638 Bytes
5c06e71 6011839 5c06e71 6011839 3db6b2b 70edc47 56c7a27 6011839 70edc47 3db6b2b 6011839 3db6b2b 68c35ce 70edc47 6011839 70edc47 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | ---
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.
|