M4A1TasteGood/ConvNeXt_Masked_Diffusion

Beyond ViT Tokens: Masked-Diffusion Pretrained Convolutional Pathology Foundation Model for Cell-Level Dense Prediction

Released weights for CMD (ConvNeXt Masked-Diffusion), a convolutional pathology foundation model for cell-level dense prediction. This repository contains only the checkpoints. Pretraining, downstream training, and inference code are in the GitHub repository.

NeurIPS 2026. Paper: arXiv:2605.08276. Code: M4A1TasteGood/ConvNeXt-Masked-Diffusion.

Checkpoints

File	What it is
weights/CMD-L/pytorch_model.bin	Pixel-space CMD-L, conditioned on H0-mini (768-d)
weights/TNBC_SegHead/best_model.pth	TNBC segmentation head trained on frozen CMD-L features

The TNBC head expects CMD-L features at image size 256 and diffusion step 50. H0-mini is not included. Download it separately and keep config.json next to pytorch_model.bin.

hf download M4A1TasteGood/ConvNeXt_Masked_Diffusion \
  weights/CMD-L/pytorch_model.bin \
  weights/TNBC_SegHead/best_model.pth \
  --local-dir .

Inference

Clone the code repository, then point it at these files:

git clone https://github.com/M4A1TasteGood/ConvNeXt-Masked-Diffusion
cd ConvNeXt-Masked-Diffusion
pip install -r requirements.txt

python infer.py \
  --image assets/test.png \
  --cmd /path/to/weights/CMD-L \
  --seg-head /path/to/weights/TNBC_SegHead/best_model.pth \
  --pathology /path/to/H0-mini/pytorch_model.bin \
  --output-dir outputs/tnbc

--cmd can be the pytorch_model.bin file or the directory that contains it. Downstream fine-tuning uses the same checkpoint through configs/downstream.yaml. See the GitHub README for the dataset format and the pretraining command.

Citation

@article{chen2026beyond,
  title={Beyond ViT Tokens: Masked-Diffusion Pretrained Convolutional Pathology Foundation Model for Cell-Level Dense Prediction},
  author={Chen, Weiming and Ling, Xitong and Cai, Zhenyang and Wang, Xidong and Li, Jiawen and Guan, Tian and Wang, Benyou and He, Yonghong},
  journal={arXiv preprint arXiv:2605.08276},
  year={2026}
}