--- license: cc-by-nc-4.0 library_name: pytorch pipeline_tag: unconditional-image-generation datasets: - data-archetype/imagenet_1k_bucketed_captioned tags: - image-generation - class-conditional - flow-matching - diffusion-transformer - pixel-space - imagenet - pytorch --- # data-archetype/pixeldit_s32_205m **pixeldit_s32_205m** is a 205M-parameter class-conditional ImageNet-1k generator that works directly in pixel space: no autoencoder runs at inference. It denoises the RGB pixels together with the 32 semantic channels of [dinac3_96][dinac3], then keeps only the pixels. > **Research preview.** pixeldit_s32_205m shares current results of a direction we > find promising: pixel-space generation guided by a few co-generated semantic > channels. It is not fully optimized: training is short, and the loss > weighting and timestep schedule have not been tuned to convergence (see the > technical report's open directions). - **State:** per 16 × 16 patch, the 768 RGB values (scaled to unit noise std) and 32 semantic channels, one 800-channel token, all denoised at one shared timestep. - **Transformer:** 16 layers at width 896, class conditioning through cross-attention to four learned tokens per class, 4 register tokens, and SPRINT token routing during training. - **Heads:** a linear readout predicts the semantic velocity; a small two-stage convolutional head (stride 8, then stride 4) that sees the noisy pixels again predicts the pixel velocity. - **Training:** flow matching on ImageNet-1k at 256 and 384 px areas in several aspect ratios, batch 256, on one GPU. **Version 0.3 (research preview)** · **[Technical report](TECHNICAL_REPORT.md)** The released weights are the EMA at training step 240,000 (batch 256, about 50 epochs of ImageNet-1k).
Reflex camera
Reflex camera		 Siberian husky
Siberian husky
Promontory
Promontory		 Jacamar
Jacamar
West Highland white terrier
West Highland white terrier		 Moped
Moped
*240k steps at batch size 256 (≈50 epochs) checkpoint, default sampler at 416 × 320.* ## Generation benchmark FID, MIND and Monge-DINO (lower is better) on our mixed-resolution ImageNet benchmark, against DiTs trained on the [dinac3_96][dinac3] and [semantic_vae][semantic-vae] latents; protocol and sampler in the [technical report, section 7.1](TECHNICAL_REPORT.md#71-generation-benchmark). | model | steps | epochs | FID ↓ | MIND ↓ | Monge-DINO ↓ | | --- | ---: | ---: | ---: | ---: | ---: | | DiT on dinac3_96 latents | 100k | ≈20.8 | 9.04 | 8.29 | 10.92 | | DiT on semantic_vae latents | 100k | ≈20.8 | 10.02 | 7.42 | 15.94 | | pixeldit_s32_205m | 100k | ≈20.8 | 10.41 | 8.14 | 22.82 | | pixeldit_s32_205m | 240k | ≈50 | 8.64 | 6.06 | 15.46 | ## Usage Download the repository and install the package, in an environment with a CUDA-enabled PyTorch (dependencies: `torch>=2.13`, `safetensors`, `huggingface-hub`, `numpy`, `scipy`, `pillow`): ```bash hf download data-archetype/pixeldit_s32_205m --local-dir pixeldit_s32_205m pip install ./pixeldit_s32_205m # optional, for grn_kernel=GrnKernel.TRITON (Linux, CUDA): pip install "./pixeldit_s32_205m[triton]" ``` ```python import torch from pixeldit_s32_205m import ( GenerationConfig, GuidanceTerm, ImageNetClass, PdgPath, PixelDiTS32Pipeline, Schedule, Solver, ) pipeline = PixelDiTS32Pipeline.from_pretrained( "data-archetype/pixeldit_s32_205m", # or a local directory device="cuda", dtype=torch.bfloat16, ) config = GenerationConfig( # the package defaults, spelled out height=320, width=416, steps=50, solver=Solver.EULER, schedule=Schedule.BETA, log_snr_shift=-2.0, cfg=None, pdg=GuidanceTerm(scale=3.0), pdg_path=PdgPath.SKIP_MIDDLE, seed=0, ) classes = [ ImageNetClass.GOLDEN_RETRIEVER, ImageNetClass.from_name("volcano"), 963, # pizza, by class index ] output = pipeline.generate(classes, config) output.to_pil()[0].save("golden_retriever.png") # images: [3, 320, 416] in [-1, 1] ``` - **Classes:** `ImageNetClass` has the 1000 ImageNet-1k classes in label order (`ImageNetClass.TENCH.value == 0`). `ImageNetClass.from_name` accepts a member name or any synonym of the class description, case-insensitive; a synonym shared by two classes (`crane`, `cardigan`) is an error naming both members. Plain indices in [0, 1000) also work. - **Sizes:** height and width are multiples of 16. Training covered 256- and 384-px areas in several aspect ratios, mostly 4:3 landscape: the default 416 × 320 (W × H) is the most common shape at 384 px, 288 × 224 at 256 px (half the cost per image). Portrait shapes swap the sides. - **Default sampler:** Euler on Beta(0.6, 0.6) quantiles, log-SNR shift −2, PDG 3 against the skipped middle layers, no CFG, 50 steps (100 model calls). Higher PDG strengthens class features but saturates colours. - **Guidance:** - `pdg` (path-drop guidance) extrapolates away from a weakened model, either with the 12 middle layers skipped (`PdgPath.SKIP_MIDDLE`) or run on one token per 2 × 2 cell (`PdgPath.THREE_QUARTER`). - `cfg` is classifier-free guidance against the null class. - Each takes a scale and an inclusive solver-step window (`GuidanceTerm(scale, start_step=0, stop_step=None)`); with both active, the update averages the two (formulas in `GuidedVelocity`). - **Solvers:** `Solver.DPMPP_2M` and `Solver.EULER`. **Schedules:** `Schedule.LINEAR`, `Schedule.BETA` (Beta(0.6, 0.6) quantiles), `Schedule.BETA_1_25`. `log_snr_shift` shifts the schedule's log-SNR (negative values spend more steps at high noise). - **Determinism:** `seed` fixes the initial noise and the three-quarter routing; `generate(..., noise=...)` takes an explicit FP32 noise tensor `[B, 800, H / 16, W / 16]`. - **Lower level:** `pipeline.model(state, t, labels, path=..., generator=...)` returns the FP32 velocity of an 800-channel state (`PixelDiTS32Model`). `from_pretrained(path_or_repo, *, device, dtype=torch.bfloat16, compile_model=True, grn_kernel=GrnKernel.REFERENCE, revision=None)`: - A `Path` is always a local directory; a `str` is a local directory if one exists, an error if it looks like a path (starts with `.`, `/` or `~`, or has more than one `/`) but names none, and a Hub repository id otherwise. Only `config.json` and `model.safetensors` are downloaded. - Inference is CUDA only. - `grn_kernel` selects how the pixel tail's FCDM blocks compute pwconv1's bias, the GELU and the GRN: `GrnKernel.REFERENCE` (default) uses plain PyTorch ops and needs nothing else, on any platform; `GrnKernel.TRITON` runs them as one fused Triton kernel (needs the `triton` extra and a CUDA GPU; loaded only when selected). The two agree to bf16 rounding. ## Precision - Set the dtype only through `from_pretrained(dtype=...)`; dtype casts after loading raise a `TypeError` (`.to(device)` works). - `torch.bfloat16` (default) runs under bf16 autocast with FP32 normalization, modulation, gating and RoPE arithmetic. With the reference GRN math it was bit-identical to the training code: from the same noise, the package's 50-step samples equalled our training repository's sampler output exactly (eager mode). The repository has since moved to the fused GRN kernel: `GrnKernel.TRITON` reproduces its velocity bit for bit, and the reference math differs from it by bf16 rounding in the pixel tail. - `torch.float32` upcasts every tensor and runs without autocast. Its images differ from bf16 by 34–43 dB PSNR after 50 steps. - The three model stages are compiled by default with dynamic shapes; the first calls compile. Compiled and eager images differ by 34–44 dB PSNR. Pass `compile_model=False` for eager. ## License - **Code** (the `pixeldit_s32_205m` package): Copyright 2026 data-archetype, licensed under the Apache License 2.0, see [LICENSE](LICENSE). - **Weights:** Creative Commons Attribution-NonCommercial 4.0 (CC BY-NC 4.0), see [LICENSE-WEIGHTS.md](LICENSE-WEIGHTS.md): research and other non-commercial use. The model was trained on ImageNet-1k, whose terms allow non-commercial research and educational use only. The `license: cc-by-nc-4.0` tag above refers to the weights. See [ATTRIBUTION.md](ATTRIBUTION.md) for sources and design credits. ## Citation ```bibtex @misc{pixeldit_s32_205m, title = {pixeldit_s32_205m: class-conditional pixel-space generation with co-generated semantic channels}, author = {data-archetype}, email = {data-archetype@proton.me}, year = {2026}, note = {Version 0.3}, url = {https://huggingface.co/data-archetype/pixeldit_s32_205m}, } ``` [dinac3]: https://huggingface.co/data-archetype/dinac3_96 [semantic-vae]: https://huggingface.co/well9472/semantic_vae