https://github.com/tonyduan/diffusion

From-scratch diffusion model implemented in PyTorch.
https://github.com/tonyduan/diffusion

diffusion generative-models pytorch

Last synced: 3 months ago
JSON representation

From-scratch diffusion model implemented in PyTorch.

• Host: GitHub
• URL: https://github.com/tonyduan/diffusion
• Owner: tonyduan
• Created: 2023-01-30T07:13:44.000Z (over 1 year ago)
• Default Branch: main
• Last Pushed: 2024-02-06T09:26:42.000Z (5 months ago)
• Last Synced: 2024-02-06T10:35:54.698Z (5 months ago)
• Topics: diffusion, generative-models, pytorch
• Language: Python
• Homepage:
• Size: 2.3 MB
• Stars: 14
• Watchers: 2
• Forks: 1
• Open Issues: 0
• Metadata Files:
  • Readme: README.md

Lists

• awesome-stars - tonyduan/diffusion - From-scratch diffusion model implemented in PyTorch. (Python)

README

        
### Diffusion Models From Scratch

March 2023.

---

My notes for this repository ended up longer than expected, too long to be rendered by GitHub.

So instead of putting notes here, they've been moved to my website.

[[**This blog post**]](https://www.tonyduan.com/diffusion/index.html) explains the intuition and derivations behind diffusion.

---

This codebase provides a *minimalist* re-production of the MNIST example below.

It clocks in at well under 500 LOC.







(Left: MNIST groundtruth. Right: MNIST sampling starting from random noise).

---

**Example Usage**

Code below is copied from `examples/ex_mnist_simple.py`, omitting boilerplate training code.

```python

    # Initialization

    nn_module = UNet(in_dim=1, embed_dim=128, dim_scales=(1, 2, 4, 8))

    model = DiffusionModel(

        nn_module=nn_module,

        input_shape=(1, 32, 32,),

        config=DiffusionModelConfig(

            num_timesteps=500,

            target_type="pred_x_0",

            gamma_type="ddim",

            noise_schedule_type="cosine",

        ),

    )

    # Training Loop

    for i in range(args.iterations):

        loss = model.loss(x_batch).mean()

        loss.backward()

    # Sampling, the number of timesteps can be less than T to accelerate

    samples = model.sample(bsz=64, num_sampling_timesteps=None, device="cuda")

```