deepsky-128px

A from-scratch denoising diffusion probabilistic model (DDPM) that unconditionally generates 128×128 deep-sky astronomical images — nebulae, galaxies, and star clusters. Every pixel is generated; nothing is retrieved or composited.

Everything is hand-written PyTorch — the U-Net, the noise schedules, the forward process, the training loop, and both samplers. No diffusers dependency.

• ▶ Live demo: https://huggingface.co/spaces/jessholbrook/deepsky
• Code + full write-up: https://github.com/jessholbrook/deepsky

Random, uncurated samples — the honest output distribution, not a hand-picked best-of.

Model description

• Type: unconditional DDPM (ε-prediction), 128×128 RGB
• Architecture: ADM-style U-Net, ~87M parameters — residual blocks with per-block timestep-embedding injection, self-attention at 16×16 and 8×8, sinusoidal timestep embeddings, zero-initialized output head, and nearest-neighbor upsampling + conv (no transposed-conv checkerboarding).
• Diffusion: Nichol & Dhariwal cosine ᾱ schedule, 1000 timesteps, derived quantities precomputed in float64.
• Samplers: DDPM ancestral (stochastic, 1000 steps) and DDIM (deterministic, η=0, arbitrary step count).

Files

File	Description
ckpt_0200000.pt	Training checkpoint (1.39 GB). EMA weights live at state["ema"]["shadow"].
samples_0200000.png	DDIM (100-step) sample grid rendered during training.
final.png	DDPM (1000-step) ancestral sample grid.

Usage

The model is a plain nn.Module from the deepsky package (no from_pretrained):

import torch
from huggingface_hub import hf_hub_download

from deepsky.config import ModelCfg
from deepsky.diffusion.gaussian import GaussianDiffusion
from deepsky.diffusion.samplers import ddim_sample
from deepsky.diffusion.schedule import make_schedule
from deepsky.models.unet import UNet

cfg = ModelCfg(
    base_channels=128, channel_mults=(1, 2, 3, 4), num_res_blocks=2,
    attn_resolutions=(16, 8), time_emb_dim=512,
)
model = UNet(cfg, 128)
state = torch.load(hf_hub_download("jessholbrook/deepsky-128px", "ckpt_0200000.pt"),
                   map_location="cpu")
model.load_state_dict(state["ema"]["shadow"])  # EMA weights = best quality
model.eval()

diffusion = GaussianDiffusion(make_schedule("cosine", 1000))
img = ddim_sample(model, diffusion, (1, 3, 128, 128), torch.device("cpu"), steps=50)

Training

• Data: 159,462 curated 256px crops (3.8 GB) from public ESA/Hubble, ESA/Webb, ESO, and NASA imagery. Heavy multi-stage curation (title/AVM-type blacklists, perceptual-hash dedup, brightness/spectra pruning, a CLIP zero-shot pass) removes charts, diagrams, and duplicates to near-zero contamination. Augmented at train time with the full dihedral group (a free 8× from rotations + flips, all physically valid for deep-sky images).
• Objective: simple ε-prediction MSE at uniformly sampled timesteps.
• Run: 200,000 steps, batch size 32, single RTX 4090, 22 h (2.48 it/s), ~$18. Adam, linear warmup + cosine LR decay 1e-4 → 1e-5, EMA 0.9999, gradient clipping, bf16 autocast. Final loss ≈ 0.01–0.05.

The model learns genuine astronomical structure on its own: point stars sharpen into crisp cores, dust lanes cut across edge-on galaxies, globular clusters resolve into thousands of colored points, and bright stars even grow the four-pointed diffraction spikes characteristic of Hubble's optics — learned purely from data, never hard-coded.

Intended use & limitations

Intended for research, education, and art: exploring how a small, from-scratch diffusion model behaves, and generating decorative astronomical imagery.

Not for scientific use. Outputs are hallucinated — they resemble real deep-sky objects statistically but depict nothing real. They must not be used as astronomical data, measurements, or evidence.

Limitations:

• Fixed 128×128 resolution; unconditional (no text or class control).
• The output distribution mirrors the curated dataset's aesthetic (Hubble/Webb/ ESO/NASA press imagery), so it inherits that visual bias.
• A minority of samples are flat, noisy, or structurally incoherent — expected at this scale and step count.

License & attribution

Weights and samples are released under CC BY 4.0. Training data derives from:

• ESA/Hubble (esahubble.org) — CC BY 4.0
• ESA/Webb (esawebb.org) — CC BY 4.0
• ESO (eso.org) — CC BY 4.0
• NASA Image and Video Library (images.nasa.gov) — public domain

Per-image credits are preserved in data/crops256/manifest.csv in the code repository. If you publish samples or weights, please retain this attribution.