Right now, datasets have to be loaded into RAM to be used in a Data object, which is an issue when dealing with very large datasets.

Add a documentation page explaining how standard image generation practices with Diffusion Models translate to Modular Diffusion.

The library currently only works for Python 3.10+ which might pose an inconvenience to some users.

As the library grows, it is important to make sure that nothing is broken when new features are added. After doing some research, I think a good way of doing this would be to create a Pytest class for each prebuilt module with a test function for each of the module's methods. As an example, I wrote this test for the Normal distribution module:

class TestNormal:

    @pytest.fixture(params=[(2, 2), (3, 3, 3), (4, 4, 4, 4)])
    def distribution(self, request: FixtureRequest) -> Normal:
        size = request.param
        mu, sigma = torch.zeros(size), torch.ones(size)
        return Normal(mu, sigma)

    def test_sample(self, distribution: Normal) -> None:
        z, epsilon = distribution.sample()
        assert z.shape == distribution.mu.shape
        assert epsilon.shape == distribution.mu.shape
        assert torch.allclose(z, distribution.mu + distribution.sigma * epsilon)

    @pytest.mark.parametrize("x", [0.0, 1.0, -1.0, 1e6])
    def test_nll(self, distribution: Normal, x: float) -> None:
        nll = distribution.nll(torch.full(distribution.mu.shape, x))
        assert nll.shape == distribution.mu.shape
        assert nll.shape == distribution.sigma.shape
        assert torch.allclose(
            nll, 0.5 * ((x - distribution.mu) / distribution.sigma)**2 +
            (distribution.sigma * 2.5066282746310002).log())

    @pytest.mark.parametrize("mu, sigma", [(0.0, 1.0), (-1.0, 3.0), (1e6, 2e6)])
    def test_dkl(self, distribution: Normal, mu: float, sigma: float) -> None:
        other = Normal(
            torch.full(distribution.mu.shape, mu),
            torch.full(distribution.sigma.shape, sigma),
        )
        dkl = distribution.dkl(other)
        assert dkl.shape == distribution.mu.shape
        assert torch.allclose(
            dkl,
            torch.log(other.sigma / distribution.sigma) +
            (distribution.sigma**2 +
             (distribution.mu - other.mu)**2) / (2 * other.sigma**2) - 0.5,
        )

I have very little experience with unit testing, so I appreciate all the help I can get with this.

How to run the Unconditional Diffusion example on custom dataset without grayscale? It returns error whenever the dataset is not grayscaled

It'd be nice to have a dark button switch at the top for those more sensitive to bright lights.

As the documentation states, the current implementations are neither the most effective or efficient. The U-Net implementation was adapted from the The Annotated Diffusion Model and the Transformer implementation was adapted from Peebles & Xie (2022) (adaptive layer norm block). Although these produce good enough results, ideally the library would provide the best implementations out there for general use.

From what I've read, I think a good choice for the U-Net implementation would be the one used in Imagen for the Text-to-Image model, but there may well be other more recent architectures that would be a better fit. For the Transformer I'm really not sure right now. Any input on this would be greatly appreciated.

The pages in question are:

• https://cabralpinto.github.io/modular-diffusion/guides/custom-modules/
• https://cabralpinto.github.io/modular-diffusion/modules/noise-type/
• https://cabralpinto.github.io/modular-diffusion/modules/loss-function/

The issue is caused by LaTeX expressions that don't fit smaller screens. I'm not quite sure how this could be fixed.

Not sure why the definition def approximate(self, z: Tensor, t: Tensor, hat: Tensor) in Gaussian noise does not consider two different parameters x and epsilon. The forged approximate distribution is the case with epsilon parameter.

Could you please have a look?

Add a documentation page explaining how standard text generation practices with Diffusion Models translate to Modular Diffusion.

Currently, we're using default Astro settings (Shiki with the github-dark theme). However, the syntax highlighting is clearly lacking, with many words not being highlighted. Useful link.

Currently, ClassifierFree guidance is hardcoded in the Model class and the Guidance base class cannot be extended to create custom guidance behavior.