FiTv2-3B-2-256/README.md

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

# FiTv2-3B-2-256

Self-contained Diffusers checkpoint for **FiTv2-3B/2**, converted from [`InfImagine/FiT`](https://huggingface.co/InfImagine/FiT).

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

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

## Recommended inference (256×256)

| Setting | Value |
| --- | --- |
| Resolution | 256×256 |
| Sampler | flow matching (velocity ODE) |
| Steps | 250 |
| CFG scale | 1.5 |
| Dtype | `float32` (or `bfloat16` on Ampere+) |
| VAE | `stabilityai/sd-vae-ft-ema` (bundled under `vae/`) |

```python
from pathlib import Path
import torch
from diffusers import DiffusionPipeline

model_dir = Path("./FiTv2-3B-2-256").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=250,
    guidance_scale=1.5,
    generator=generator,
).images[0]
image.save("demo.png")
```