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, 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

The released weights are the EMA at training step 240,000 (batch 256, about 50 epochs of ImageNet-1k).

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 and semantic_vae latents; protocol and sampler in the technical report, section 7.1.

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):

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]"

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.
• Weights: Creative Commons Attribution-NonCommercial 4.0 (CC BY-NC 4.0), see 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 for sources and design credits.

Citation

@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},
}