Self-Flow, Black Forest Labs inference code for an ImageNet 256 diffusion model

What is actually in the tree: one sampling script and three modules

The README is unusually direct about scope. The folder contains inference code for generating images with the Self-Flow trained diffusion model on ImageNet 256x256, and it says the code allows you to load a checkpoint pretrained on ImageNet 256x256 and generate 50,000 images for FID evaluation.

That is the whole product. The project structure section lists a main sampling script, a checkpoints directory for model checkpoints, the requirements file, the README and a src directory:

Self-Flow/
├── sample.py           # Main sampling script
├── checkpoints/        # Place model checkpoints here
├── requirements.txt    # Python dependencies
├── README.md           # This file
└── src/                # Model and sampling implementations
    ├── model.py        # SelfFlowPerTokenDiT model
    ├── sampling.py     # Diffusion sampling utilities
    └── utils.py        # Position encoding utilities

So the model definition lives in src/model.py as a SelfFlowPerTokenDiT class, the sampler in src/sampling.py, and position encoding helpers in src/utils.py. There is no training script, no dataset preparation code and no configuration system. The repository also has no GitHub releases, so version pinning is a matter of using the commit you read.

The likely audience is therefore a paper reader or a replication attempt, not a product team. The two capabilities the code offers, sample in bulk and score the samples with the ADM evaluation suite, are exactly the two steps needed to check the reported numbers.

The licence is Apache-2.0, which is permissive and includes an express patent grant, and the code builds on two earlier projects the README names: REPA for representation alignment, and SiT for scalable interpolant transformers.

The idea: flow matching with a self-supervised reconstruction objective

Self-Flow is described as a training framework that combines the flow matching objective with a self-supervised feature reconstruction objective. The repository itself is the inference half of that framework, but the training configuration section describes enough to understand what the checkpoint represents.

The training details list a model of SiT-XL/2 with per-token timestep conditioning, Self-Flow training with per-token masking at a 25 percent mask ratio, AdamW as the optimizer with gradient clipping at max_norm of 1, and mixed precision in BFloat16. The self-distillation setup names a teacher at layer 20 maintained as an exponential moving average and a student at layer 8.

That layer split is the interesting part for anyone trying to reason about the model. A deeper teacher supplying targets to a shallower student is a distillation arrangement rather than a plain teacher-student copy, and the EMA teacher is the moving target that makes the self-supervised objective stable.

The claimed scope of the method is multi-modal synthesis, which the title of the paper carries, while the released checkpoint and this code cover the image case on ImageNet at 256x256. A reader looking for text to image or audio behaviour should not infer that from the repository.

The paper is the primary source for the method, listed as arXiv preprint 2603.06507, with a project page hosted by Black Forest Labs and a BibTeX entry included in the README for citation.

Per-token timestep conditioning in the SiT-XL/2 backbone

The architectural modification the README highlights is per-token timestep conditioning, and the explanation is short enough to be unambiguous: it allows each token to have a different noise level during training.

In a standard diffusion transformer, one timestep applies to the whole sequence, so every token in an image is denoised from the same noise level at each step. Per-token conditioning breaks that, giving each token its own noise level. For a sampler that then has to integrate a flow matching ODE or SDE with heterogeneous noise levels per token, that is a substantial change to how the network is conditioned, and it is the mechanism that connects the training objective to the inference code in src/model.py.

The base architecture is SiT-XL/2, the scalable interpolant transformer from the SiT project the README credits. Using an existing backbone rather than a new one means the checkpoint is comparable to other SiT based results, and the self-distillation and masking described above are the parts that are specific to Self-Flow.

A practical note for anyone reading the code: the per-token scheme is the reason the sampling code exposes both an SDE and an ODE mode, since a stochastic sampler and a deterministic one follow different trajectories through the same conditioned network.

Running the sampler: dependencies, checkpoint and the 50,000 sample run

The environment setup is a single requirements install:

pip install -r requirements.txt

The requirements file pins torch at 2.0.0 or newer, torchvision, diffusers, transformers and timm for the model stack, numpy, scipy and Pillow for numerical and image work, einops and tqdm, plus tensorflow at 2.10.0 or newer, which is there for the ADM evaluator rather than for sampling.

The checkpoint comes from Hugging Face, and the README gives the download as a one-liner that writes into a local checkpoints directory:

python -c "
from huggingface_hub import hf_hub_download
hf_hub_download(
    repo_id='Hila/Self-Flow',
    filename='selfflow_imagenet256.pt',
    local_dir='./checkpoints'
)
print('Downloaded!')
"

Sampling for the evaluation run is the multi-GPU path, and the README marks multi-GPU as recommended because the target is fifty thousand images:

torchrun --nnodes=1 --nproc_per_node=8 sample.py \
    --ckpt checkpoints/selfflow_imagenet256.pt \
    --output-dir ./samples \
    --num-fid-samples 50000

The single GPU path adds an explicit batch size:

python sample.py \
    --ckpt checkpoints/selfflow_imagenet256.pt \
    --output-dir ./samples \
    --num-fid-samples 50000 \
    --batch-size 64

Beyond those two flags, the argument table lists --num-steps defaulting to 250, --mode defaulting to SDE with ODE as the alternative, --seed defaulting to 31 for reproducibility, and --cfg-scale defaulting to 1.0, which the README notes means no guidance, as used in the paper. The output directory defaults to ./samples, and the sample count defaults to 50000. Those defaults reproduce the paper setting, so changing them turns your run into an ablation rather than a reproduction.

Scoring the samples with the ADM evaluation suite

The samples are written as an .npz file, and the README states that this file can be used with the ADM evaluation suite to compute FID, IS, Precision and Recall. The evaluation code itself is not in this repository, which keeps the dependency surface smaller but adds a step.

The evaluation needs reference statistics for ImageNet at 256 resolution, and the README points at the OpenAI public blob storage:

wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz

With the reference batch and the generated samples in place, the evaluator is run as:

python evaluator.py \
    VIRTUAL_imagenet256_labeled.npz \
    ./samples/samples_50000.npz ./samples

The file naming matters here. The generated archive is expected to be named samples_50000.npz inside the output directory, which ties the file name to the sample count, so generating a different number of samples changes the expected path.

Two practical warnings. The ADM suite is a tensorflow based evaluator, which is why tensorflow appears in the requirements file, and it needs the reference statistics downloaded from a third party host. Anyone reproducing the numbers should treat the FID value as meaningful only when the sample count, the reference batch, the seed and the sampler settings all match what the README describes.

Limits of the release, licence and maintenance

The first limit is that this is inference code for one checkpoint on one resolution. There is no text conditioning, no image to image path and no training entry point in the tree, so the interesting half of the research, the self-supervised objective and the distillation schedule, cannot be run from here. Reproducing the training run means writing it against the paper.

The second limit is the evaluation cost. Fifty thousand samples on a diffusion model is a substantial generation job, which is why the README recommends multi-GPU, and the scoring step adds a tensorflow based suite on top. The FID figure is a property of the whole pipeline, checkpoint plus sampler plus sample count plus reference batch, and it is easy to produce a number that is close to the paper without being comparable to it.

The third is packaging. There are no GitHub releases, so a consumer pins a commit, and the repository depends on other projects by name, REPA and SiT, whose APIs could move.

The last push to the repository was on 2026-05-23, and the project is presented as the code and website for an ICML'26 paper with a project page at bfl.ai/research/self-flow. The licence is Apache-2.0, which permits commercial use and modification with attribution and notice, and includes the patent grant that makes it a reasonable choice for a company building on top of it. That is a licensing note rather than legal advice, and anyone shipping a product should read the LICENSE file itself.

Editorial conclusion

This repository fits researchers who want to sample from the released Self-Flow checkpoint and reproduce its FID numbers, and it does not fit teams looking for a text to image model, a training pipeline or a maintained library API, because the tree contains a single sampling script plus a small src directory and no published releases. Verify first that the checkpoint file selfflow_imagenet256.pt downloads from the Hila/Self-Flow repository on Hugging Face, that the dependencies in requirements.txt, including torch and tensorflow for the ADM evaluator, install in your environment, and that the sampling defaults match the paper setting before you change --mode, --num-steps or --cfg-scale away from their shipped values.

Frequently asked questions

Official sources

1. black-forest-labs/Self-Flow on GitHub
2. Issues
3. License: Apache-2.0
4. Project website
5. README