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.