fushinguyenex/IMF

Multi-Trace U-Net — 14,992,182 parameters

U-Net adapted from the user-provided Untitled0.ipynb and widened to approximately 22M parameters. Trained for 50,000 optimizer updates on CIFAR-10. Training samples K uniformly from 1..5 and a grid interval j/K; the same u head fits the interval and diagonal, with a detached directional JVP. Default inference weights are ema. Model input/output pixels are 32x32 RGB; sampling maps the final [-1,1] image into [0,1].

This is a custom PyTorch model. Download with huggingface_hub.snapshot_download, install the included requirements.txt, then run:

python sample.py --model-dir . --nfe 1 --output samples_1nfe.png
python sample.py --model-dir . --nfe 5 --output samples_5nfe.png

config.json describes architecture/objective; training_config.json records the effective batch, accumulation, LR schedule, clipping and EMA decay. Full optimizer/EMA/RNG resume checkpoints remain on the Modal Volume and, when the Modal notebook backup option is enabled, are also stored under training-checkpoints// in this Hub repository. best_validation.safetensors contains the selected EMA validation snapshot; its step and fixed held-out losses are recorded in best_validation.json. Diagnostics include gradient RMS, actual optimizer updates, losses by K, clipping frequency, samples, throughput, ETA and TensorBoard events. Plateau flags are heuristics, not proof of convergence or image quality.

When present, fid.json reports pytorch-fid pool3 FID against 50,000 unaugmented CIFAR-10 training images, with the generated sample count and each NFE explicitly recorded. This protocol differs from the reference notebook's ImageNet-feature FID; the values must not be compared directly. No quality claim is made here. The test split was used as fixed validation during training.

Labels: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck. The U-Net and interval schedule follow the user-provided notebook; repository training flow and DiT baseline are adapted from haidog-yaqub/MeanFlow (MIT). See third_party/ for attribution.