Hackable Diffusion (PyTorch)

Hackable Diffusion is a modular, composable, hackable library for generative diffusion models in PyTorch.

Ported from Google DeepMind's google/hackable_diffusion (originally written in JAX/Flax), this library brings clean abstractions and high composability to PyTorch users across continuous Gaussian diffusion, discrete masked diffusion, simplicial Dirichlet probability diffusion, Riemannian manifold flow matching, and multi-modal joint PyTree diffusion.

---

🌟 Key Design Goals & Advantages in PyTorch

• 🎨 Hackable & Modular: Every corruption schedule, timestep sampler, neural architecture, guidance function, and step sampler is an independent component that can be swapped or customized.
• ⚡ PyTorch-Native Elegance: Replaces JAX functional pure-function patterns (such as explicit state dictionaries and PRNGKey passing) with standard torch.nn.Module, .train() / .eval() states, and native PyTorch RNGs.
• 🌳 PyTree Composability: Native support for PyTorch dict, tuple, and list data structures without requiring external PyTree registration.
• 🚀 Performance Optimized: Uses hardware-accelerated F.scaled_dot_product_attention (SDPA / FlashAttention) where available.

---

🛠️ Installation

# Clone the repository
git clone https://github.com/monatis/hackable-diffusion.git
cd hackable-diffusion

# Install dependencies using uv
uv sync

---

🚀 Quickstart Examples

1. Continuous Flow Matching & Diffusion (toy_2d_demo.py)

import torch
from hackable_diffusion.corruption import GaussianProcess, RFSchedule
from hackable_diffusion.training import UniformTimeSampler, NoWeightGaussianLoss
from hackable_diffusion.sampling import DiffusionSampler, UniformTimeSchedule, EulerFlowMatchingStep
from hackable_diffusion.inference import PyTorchInferenceFn

# Process & Loss setup
schedule = RFSchedule()  # Rectified Flow (alpha=1-t, sigma=t)
process = GaussianProcess(schedule=schedule)
time_sampler = UniformTimeSampler()
loss_fn = NoWeightGaussianLoss(prediction_type="velocity")

# Model training loop snippet
x0 = torch.randn(64, 2)
time = time_sampler(shape=(64,))
xt, targets = process.corrupt(None, x0, time)
preds = model(time=time, xt=xt, is_training=True)
loss = loss_fn(preds, targets, time).mean()

# Sampling loop
inference_fn = PyTorchInferenceFn(network=model)
sampler = DiffusionSampler(
    time_schedule=UniformTimeSchedule(),
    stepper=EulerFlowMatchingStep(process=process),
    num_steps=50,
)
init_noise = torch.randn(16, 2)
final_step, trajectory = sampler(inference_fn=inference_fn, rng=None, initial_noise=init_noise)

2. Discrete Masked Diffusion (mnist_discrete_demo.py)

from hackable_diffusion.corruption import CategoricalProcess, LinearDiscreteSchedule
from hackable_diffusion.training import MD4Loss
from hackable_diffusion.sampling import MaskedStepSampler

process = CategoricalProcess.masking_process(
    schedule=LinearDiscreteSchedule(),
    num_categories=10,
)
loss_fn = MD4Loss(schedule=process.schedule)
stepper = MaskedStepSampler(process=process)

3. Simplicial Dirichlet Diffusion (mnist_simplicial_demo.py)

from hackable_diffusion.corruption import SimplicialProcess, LinearDiscreteSchedule
from hackable_diffusion.sampling import SimplicialStepSampler

process = SimplicialProcess.uniform_process(
    schedule=LinearDiscreteSchedule(),
    num_categories=5,
    temperature=1.0,
)
stepper = SimplicialStepSampler(process=process)

4. Riemannian Flow Matching on Manifolds (riemannian_sphere_demo.py)

from hackable_diffusion.manifolds import Sphere
from hackable_diffusion.architecture.riemannian import RiemannianConditionalBackbone
from hackable_diffusion.sampling import RiemannianEulerStep

sphere = Sphere(dim=2) # S^2 in R^3
riemannian_backbone = RiemannianConditionalBackbone(base_backbone=base_mlp, manifold=sphere)
stepper = RiemannianEulerStep(manifold=sphere)

5. Multi-Modal Joint PyTree Diffusion (multimodal_demo.py)

from hackable_diffusion.multimodal import NestedProcess, NestedSamplerStep

nested_process = NestedProcess(processes={
    "image": image_process,
    "depth": depth_process,
})
nested_stepper = NestedSamplerStep(steppers={
    "image": EulerFlowMatchingStep(process=image_process),
    "depth": EulerFlowMatchingStep(process=depth_process),
})

6. Block-Causal Discrete Diffusion LLMs (qwen_block_diffusion_demo.py)

from hackable_diffusion.llm import (
    BlockCausalDiffusionModel,
    BlockDiscreteStepSampler,
    corrupt_canvas_tokens,
)

model = BlockCausalDiffusionModel(
    base_model=base_qwen_or_llama,
    block_size=32,
    self_cond_rate=0.5,
)
stepper = BlockDiscreteStepSampler(process=process, block_size=32, prompt_len=prompt_len)

---

📦 Package Architecture Overview

src/hackable_diffusion/
├── hd_api.py                 # Protocol interfaces (CorruptionProcess, SamplerStep, etc.)
├── hd_typing.py              # Core type aliases & TargetInfo dict definitions
├── pytorch_helpers.py        # PyTree mapping, broadcasting & SafeSpan helpers
├── manifolds.py              # Euclidean, Sphere (S^n), SO(3) Lie group geometry
├── fast_random.py            # Vectorized Gumbel-Max & logit sampling utilities
├── architecture/             # Modular UNet, DiT, MLP, Attention & Norm blocks
├── corruption/               # Continuous Gaussian, Discrete & Simplicial noise processes
├── training/                 # Loss functions & timestep sampling strategies
├── sampling/                 # Continuous, Discrete, Simplicial & Riemannian step samplers
├── inference/                # Guidance rules & dynamic thresholding projections
├── multimodal.py             # PyTree wrappers for joint multi-modal diffusion
└── llm/                      # Block-Causal Discrete Diffusion LLMs (DiffusionGemma style)

---

🧪 Testing & Runnable Demos

Run the complete PyTorch test suite:

uv run pytest tests/

Run any of the 8 runnable demo scripts in examples/:

uv run examples/toy_2d_demo.py
uv run examples/mnist_unet_demo.py
uv run examples/mnist_dit_demo.py
uv run examples/mnist_discrete_demo.py
uv run examples/mnist_simplicial_demo.py
uv run examples/riemannian_sphere_demo.py
uv run examples/multimodal_demo.py
uv run examples/qwen_block_diffusion_demo.py

---

📜 License

Apache License 2.0. See LICENSE for details.