shy0423/JiT-B-FD-SIM-bs1024

JiT-B · FD-SIM post-trained (bs1024) · 1-step ImageNet-256

对 JiT-B 基座做 FD-SIM 后训练得到的 单步（1-NFE） ImageNet-256 条件生成权重。

• 损失：FD-SIM = SigLIP + MAE + Inception 三个表征空间上的单高斯 Bures–Wasserstein Fréchet Distance（fid/(fid+ε) 归一化）
• 训练：global batch 1024，cfg=1.0，num_sampling_steps=1，共 62500 step
• 发布权重：每个 milestone 取该步 best-EMA（评测最优）写入 model，格式与官方 JiT-B_FD-SIM.pth 一致，可直接 --load_from

代码：shihaoyang0423/FD-Loss（基于 Jiawei-Yang/FD-Loss）。上游基座与参考统计量见 jjiaweiyang/FD-Loss。

Files

checkpoints/
  JiT-B_FD-SIM_bs1024_step12500.pth   # ema=edm_500
  JiT-B_FD-SIM_bs1024_step25000.pth   # ema=edm_500
  JiT-B_FD-SIM_bs1024_step37500.pth   # ema=edm_500
  JiT-B_FD-SIM_bs1024_step50000.pth   # ema=edm_500
  JiT-B_FD-SIM_bs1024_step62500.pth   # ema=edm_250  (final)

每个文件约 0.53 GB，键：model / step / samples_seen / ema_label。

Results (1-step, cfg=1.0, 50k)

step	EMA	FID(ADM) ↓	FDr⁶ ↓
12.5k	edm_500	3.298	7.609
25k	edm_500	1.987	6.103
37.5k	edm_500	1.311	5.388
50k	edm_500	1.027	5.237
62.5k	edm_250	0.944	5.063

Download

# 中国大陆建议走镜像
export HF_ENDPOINT=https://hf-mirror.com

pip install -U huggingface_hub
hf download shy0423/JiT-B-FD-SIM-bs1024 \
  --local-dir . \
  --include "checkpoints/*.pth"

Sampling（1-step）

依赖本仓库代码（FD-Loss）。核心是 JiTDenoiser.generate：从噪声 t=1 欧拉一步到 t=0。

评测 / 批量出图（推荐）

export HF_ENDPOINT=https://hf-mirror.com   # 如需下依赖权重

CKPT=checkpoints/JiT-B_FD-SIM_bs1024_step62500.pth

torchrun --nproc_per_node=8 eval_all_fds.py \
  --model JiT_B \
  --rope_2d --learned_pe --legacy_time_convention --ema_type edm \
  --cfg 1.0 --cfg_list 1.0 \
  --interval_min 0.1 --interval_max 1.0 \
  --num_sampling_steps 1 \
  --eval_ema_labels online \
  --eval_bsz 128 --num_images 50000 \
  --load_from "$CKPT" \
  --output_dir work_dirs/eval_simfull1024 \
  --project eval --exp_name JiT-B-FD-SIM-bs1024-62500

等价地也可：

PRESET=JiT_B \
CKPT_PATH=checkpoints/JiT-B_FD-SIM_bs1024_step62500.pth \
CFG_OVERRIDE=1.0 \
bash scripts/evaluate_released_ckpt.sh

发布权重已是 best-EMA，因此 --eval_ema_labels online 即可复现上表 FID。本 run 的 1-step 评测 cfg 固定为 1.0。

最小采样片段

import argparse
import torch
from torchvision.utils import save_image

from utils.builders import create_generation_model
from utils.checkpoint_util import ckpt_resume
from utils.sampling_util import generate_images

args = argparse.Namespace(
    model="JiT_B",
    img_size=256,
    num_classes=1000,
    label_drop_prob=0.1,
    attn_dropout=0.0,
    proj_dropout=0.0,
    P_mean=-0.8,
    P_std=0.8,
    t_eps=0.05,
    rope_2d=True,
    learned_pe=True,
    legacy_time_convention=True,
    ema_type="edm",
    ema_rates=None,
    ema_halflife_kimg=[250, 500, 1000, 2000],
    load_from="checkpoints/JiT-B_FD-SIM_bs1024_step62500.pth",
    resume_from=None,
    auto_resume=False,
    num_sampling_steps=1,      # 一步生成
    sampling_method="euler",
    cfg=1.0,
    interval_min=0.1,
    interval_max=1.0,
    same_noise=False,
    enable_amp=True,
    amp_dtype=torch.bfloat16,
    # builders 里其它默认字段按 main_fd / eval_all_fds 的 parser 补齐即可
)

model, ema_model = create_generation_model(args)
ckpt_resume(args, model, optimizer=None, model_ema=ema_model)

labels = torch.randint(0, 1000, (8,), device="cuda")
images = generate_images(args, model, labels=labels, cfg=1.0)  # [0,1]
save_image(images, "samples.png", nrow=4)

generate_images → JiTDenoiser.generate：

1. z ~ N(0, I)（noise_scale=1）
2. t: 1 → 0，num_sampling_steps=1 时只做一次 Euler：z ← z + (0-1)·v_θ(z,t=1,y)
3. 输出像素空间 [-1,1]，再线性映射到 [0,1]

训练时的可微 1-step 路径（sample_images_with_grad）把 cfg 固定为 1.0；评测/采样走 generate，本发布建议 cfg=1.0。

Citation

若使用本权重，请同时引用 FD-Loss：

@article{yang2026fdloss,
  title={Representation Fr\'echet Loss for Visual Generation},
  author={Yang, Jiawei and Geng, Zhengyang and Ju, Xuan and Tian, Yonglong and Wang, Yue},
  journal={arXiv:2604.28190},
  year={2026}
}

License

MIT（与上游 FD-Loss 一致）。