---
license: apache-2.0
library_name: diffusers
pipeline_tag: unconditional-image-generation
tags:
  - diffusers
  - dit-moe
  - image-generation
  - class-conditional
  - imagenet
inference: true
---

# DiT-MoE-S-8E2A

Self-contained Diffusers checkpoint for **DiT-S/2** with MoE routing, converted from [`feizhengcong/DiT-MoE`](https://huggingface.co/feizhengcong/DiT-MoE).

Each subfolder is a self-contained Diffusers model repo with:

- `model_index.json` (includes ImageNet `id2label`)
- `pipeline.py` (custom `DiTMoEPipeline`)
- `transformer/transformer_dit_moe.py` and weights
- `vae/diffusion_pytorch_model.safetensors`
- `scheduler/scheduler_config.json`

## ImageNet class labels

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

- `pipe.id2label[207]` — `"golden retriever"`
- `pipe.get_label_ids("golden retriever")` — `[207]`
- `pipe(class_labels="golden retriever", ...)` — string labels resolved automatically

## Recommended inference (256×256)

| Setting | Value |
| --- | --- |
| Resolution | 256×256 |
| Sampler | DDIM |
| Steps | 50 |
| CFG scale | 4.0 |
| Dtype | `bfloat16` (recommended on Ampere+) |
| VAE | `stabilityai/sd-vae-ft-mse` (bundled under `vae/`) |

```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 = 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")
```