ModelListedExternal

DiT MoE Diffusers

Diffusers implementation of DiT-MoE (Diffusion Transformer with Mixture of Experts) for class-conditional ImageNet generation. Each variant folder is self-contained:

Sourcehuggingface.co/BiliSakura/DiT-MoE-diffusers/tree/343b8d8ef8ef83a5258db9af9abc2331dc27f6bd

KindOtherVersionBgrade

Model source

Kind: Other
File size: 319 MB
Version: vdiffusion_pytorch_model
Updated: Oct 9, 2026
Source: Hugging Face

Source description

DiT-MoE-diffusers

Diffusers implementation of DiT-MoE (Diffusion Transformer with Mixture of Experts) for class-conditional ImageNet generation. Each variant folder is self-contained:

pipeline.py — DiTMoEPipeline
scheduler/scheduler_config.json — DDIMScheduler (S/B) or DiTMoEFlowMatchScheduler (XL/G)
transformer/transformer_dit_moe.py — DiTMoETransformer2DModel
vae/ — AutoencoderKL (stabilityai/sd-vae-ft-mse)

ImageNet class labels

Each variant keeps an English id2label map directly in its own model_index.json (DiT-style).

pipe.id2label — inspect id → English label correspondence
pipe.labels — reverse map (English synonym → id), sorted for browsing
pipe.get_label_ids("golden retriever")
pipe(class_labels="golden retriever", ...) — string labels resolved automatically

Available checkpoints

Checkpoint	Path	Resolution	Sampler
DiT-MoE-S/2-8E2A	`./DiT-MoE-S-8E2A`	256×256	DDIM
DiT-MoE-B/2-8E2A	`./DiT-MoE-B-8E2A`	256×256	DDIM
DiT-MoE-XL/2-8E2A	`./DiT-MoE-XL-8E2A`	256×256	RF
DiT-MoE-G/2-16E2A	(convert with `--rectified-flow --num-experts 16`)	256×256	RF

Convert from official weights

conda activate rsgen
cd libs/DiT-MoE-diffusers

python scripts/convert_dit_moe_to_diffusers.py \
  --checkpoint ../../models/feizhengcong/DiT-MoE/dit_moe_s_8E2A.pt \
  --output ../../models/BiliSakura/DiT-MoE-diffusers/DiT-MoE-S-8E2A \
  --model DiT-S/2 \
  --num-experts 8 \
  --num-experts-per-tok 2 \
  --copy-vae ../../models/feizhengcong/DiT-MoE/sd-vae-ft-mse \
  --check-load

Inference

Use torch.bfloat16 on Ampere+ GPUs (default in examples and sample_dit_moe.py).

from pathlib import Path
import torch
from diffusers import DiffusionPipeline

model_dir = Path("./DiT-MoE-S-8E2A").resolve()
pipe = DiffusionPipeline.from_pretrained(
    str(model_dir),
    local_files_only=True,
    custom_pipeline=str(model_dir / "pipeline.py"),
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")

print(pipe.id2label[207])
print(pipe.get_label_ids("golden retriever"))

generator = torch.Generator(device="cuda").manual_seed(42)
image = pipe(
    class_labels="golden retriever",
    height=256,
    width=256,
    num_inference_steps=50,
    guidance_scale=4.0,
    generator=generator,
).images[0]
image.save("demo.png")

Citation

@article{FeiDiTMoE2024,
  title={Scaling Diffusion Transformers to 16 Billion Parameters},
  author={Zhengcong Fei and Mingyuan Fan and Changqian Yu and Debang Li and Jusnshi Huang},
  year={2024},
  journal={arXiv preprint arXiv:2407.11633},
}

Sources

1 source

Hugging FacePrimaryVersion 343b8d8ef8ef83a5258db9af9abc2331dc27f6bd

huggingface.co/BiliSakura/DiT-MoE-diffusers/tree/343b8d8ef8ef83a5258db9af9abc2331dc27f6bd

Verified Oct 9

Download files

Model artifacts

3 artifacts

DiT-MoE-B-8E2A/vae/diffusion_pytorch_model.safetensors

safetensors · 319 MB · SHA-256 f3944a3bb4da…6bac · Hugging Face

Download file

DiT-MoE-S-8E2A/vae/diffusion_pytorch_model.safetensors

safetensors · 319 MB · SHA-256 f3944a3bb4da…6bac · Hugging Face

Download file

DiT-MoE-XL-8E2A/vae/diffusion_pytorch_model.safetensors

safetensors · 319 MB · SHA-256 f3944a3bb4da…6bac · Hugging Face

Download file

Source excerpts

2 excerpts

ReadmeAug 1huggingface.co/BiliSakura/DiT-MoE-diffusers

--- license: apache-2.0 language: - en library_name: diffusers tags: - diffusers - image-generation - class-conditional - dit-moe pipeline_tag: unconditional-image-generation --- # DiT-MoE-diffusers Diffusers implementation of **DiT-MoE** (Diffusion Transformer with Mixture of Experts) for class-conditional ImageNet generation. Each variant folder is self-contained: - `pipeline.py` — `DiTMoEPipeline` - `scheduler/scheduler_config.json` — `DDIMScheduler` (S/B) or `DiTMoEFlowMatchScheduler` (XL/G) - `transformer/transformer_dit_moe.py` — `DiTMoETransformer2DModel` - `vae/` — `AutoencoderKL` (`stabilityai/sd-vae-ft-mse`) ## ImageNet class labels Each variant keeps an English `id2label` map directly in its own `model_index.json` (DiT-style). - `pipe.id2label` — inspect id → English label correspondence - `pipe.labels` — reverse map (English synonym → id), sorted for browsing - `pipe.get_label_ids("golden retriever")` - `pipe(class_labels="golden retriever", ...)` — string labels resolved automatically ## Available checkpoints | Checkpoint | Path | Resolution | Sampler | | --- | --- | --- | --- | | DiT-MoE-S/2-8E2A | `./DiT-MoE-S-8E2A` | 256×256 | DDIM | | DiT-MoE-B/2-8E2A | `./DiT-MoE-B-8E2A` | 256×256 | DDIM | | DiT-MoE-XL/2-8E2A | `./DiT-MoE-XL-8E2A` | 256×256 | RF | | DiT-MoE-G/2-16E2A | *(convert with `--rectified-flow --num-experts 16`)* | 256×256 | RF | ## Convert from official weights ```bash conda activate rsgen cd libs/DiT-MoE-diffusers python scripts/convert_dit_moe_to_diffusers.py \ --checkpoint ../../models/feizhengcong/DiT-MoE/dit_moe_s_8E2A.pt \ --output ../../models/BiliSakura/DiT-MoE-diffusers/DiT-MoE-S-8E2A \ --model DiT-S/2 \ --num-experts 8 \ --num-experts-per-tok 2 \ --copy-vae ../../models/feizhengcong/DiT-MoE/sd-vae-ft-mse \ --check-load ``` ## Inference Use `torch.bfloat16` on Ampere+ GPUs (default in examples and `sample_dit_moe.py`). ```python from pathlib import Path import torch from diffusers import DiffusionPipeline model_dir = Path("./DiT-MoE-S-8E2A").resolve() pipe = DiffusionPipeline.from_pretrained( str(model_dir), local_files_only=True, custom_pipeline=str(model_dir / "pipeline.py"), trust_remote_code=True, torch_dtype=torch.bfloat16, ) pipe.to("cuda") print(pipe.id2label[207]) print(pipe.get_label_ids("golden retriever")) generator = torch.Generator(device="cuda").manual_seed(42) image = pip...

NoteAug 1huggingface.co/BiliSakura/DiT-MoE-diffusers

Source context: 0 downloads · 0 likes · Pipeline unconditional-image-generation · Library diffusers · Repo BiliSakura/DiT-MoE-diffusers