dic-h-256x256-sit-sde

DiC-H 256x256 Checkpoint with SiT SDE Sampling

This repository packages the DiC-H ImageNet 256x256 checkpoint provided as:

pretrained_models/DiC_H_flow_2000K_256_conv.pt

The checkpoint is a DiC ConvUNet checkpoint trained for the SiT/flow transport setting and sampled with Saining Xie's SiT transport/SDE sampler. It is not the official SiT-XL/2 checkpoint.

Files

Checkpoint Metadata

Environment

conda env create -f environment.yml
conda activate SiT
pip install diffusers torchvision tqdm pillow numpy

The sampler downloads stabilityai/sd-vae-ft-mse by default unless it is already cached.

Single-Image SDE Sampling

This command uses the DiC-H checkpoint with SiT's SDE sampler and classifier-free guidance:

python sample_dic_sde.py \
  --model DiC-H \
  --ckpt pretrained_models/DiC_H_flow_2000K_256_conv.pt \
  --image-size 256 \
  --vae mse \
  --labels 207 \
  --seed 0 \
  --cfg-scale 4.0 \
  --num-sampling-steps 250 \
  --sampling-method Euler \
  --diffusion-form sigma \
  --diffusion-norm 1.0 \
  --last-step Mean \
  --last-step-size 0.04 \
  --path-type Linear \
  --prediction velocity

For this exact local checkpoint file, the smoke-test run used partial checkpoint recovery:

python sample_dic_sde.py \
  --model DiC-H \
  --ckpt pretrained_models/DiC_H_flow_2000K_256_conv.pt \
  --image-size 256 \
  --vae mse \
  --labels 207 \
  --seed 0 \
  --cfg-scale 4.0 \
  --num-sampling-steps 250 \
  --sampling-method Euler \
  --diffusion-form sigma \
  --diffusion-norm 1.0 \
  --last-step Mean \
  --last-step-size 0.04 \
  --path-type Linear \
  --prediction velocity \
  --allow-partial-corrupt-checkpoint

The local smoke-test output is:

samples/dic_h_sde_partial_seed0.png

FID-50K Sampling Command

Use this command for DiC-H FID-50K sampling with the same SiT SDE transport settings:

torchrun --nnodes=1 --nproc_per_node=8 --master_port=29683 sample_dic_sde_ddp.py SDE \
  --model DiC-H \
  --ckpt pretrained_models/DiC_H_flow_2000K_256_conv.pt \
  --image-size 256 \
  --vae mse \
  --sample-dir samples/dic_h_sit_sde_cfg_fid \
  --per-proc-batch-size 4 \
  --num-fid-samples 50000 \
  --cfg-scale 4.0 \
  --num-sampling-steps 250 \
  --sampling-method Euler \
  --diffusion-form sigma \
  --diffusion-norm 1.0 \
  --last-step Mean \
  --last-step-size 0.04 \
  --path-type Linear \
  --prediction velocity

The script creates a folder of PNG samples and an ADM-compatible .npz next to that folder.

Local Verification Status

The packaged code was verified against the provided local file:

python sample_dic_sde.py \
  --model DiC-H \
  --ckpt pretrained_models/DiC_H_flow_2000K_256_conv.pt \
  --image-size 256 \
  --dry-run-load \
  --allow-partial-corrupt-checkpoint \
  --device cpu \
  --num-sampling-steps 1

Result:

Partially loaded 692 tensors; skipped 15 corrupt tensors.
DiC-H parameters: 1,034,502,916
Checkpoint loaded successfully.
Dry-run load completed; skipping VAE decode and sampling.

A valid FID-50K number should be reported only from a checkpoint that loads completely with standard torch.load. The exact local checkpoint file packaged here requires partial recovery, so the included smoke-test image is only a loading/sampling sanity check and not a valid FID reproduction.

Citation

@article{tian2025dic,
  author       = {Yuchuan Tian and Jing Han and Chengcheng Wang and Yuchen Liang and Chao Xu and Hanting Chen},
  title        = {DiC: Rethinking Conv3x3 Designs in Diffusion Models},
  journal      = {CoRR},
  volume       = {abs/2501.00603},
  year         = {2025},
  doi          = {10.48550/ARXIV.2501.00603},
  eprinttype   = {arXiv},
  eprint       = {2501.00603}
}