TorchLean API

NN.API.Models.Diffusion

Diffusion Models #

Config-style diffusion model constructors plus reusable, dataset-independent DDPM/DDIM helpers.

The runnable examples decide where data comes from (CIFAR-10, ImageNet-style folders, synthetic artifacts). The definitions here are shape-parametric and can be reused by tests, examples, and future proof layer specifications.

Configuration for a minimal epsilon-predictor conv net.

Instances For

    Epsilon-predictor input shape, with one extra channel carrying the diffusion time.

    Instances For

      Epsilon-predictor output shape matching the denoised data channels.

      Instances For
        def TorchLean.nn.models.EpsConvNetConfig.sameConv {d : } (cfg : EpsConvNetConfig d) (inChannels outChannels : ) [NeZero inChannels] :

        Seeded shape-preserving convolution over an arbitrary spatial rank.

        Instances For
          def TorchLean.nn.models.epsConvNet {d : } (cfg : EpsConvNetConfig d) (h_batch : cfg.batch 0 := by decide) (h_dataC : cfg.dataChannels 0 := by decide) (h_inC : cfg.dataChannels + 1 0 := by decide) (h_hiddenC : cfg.hiddenChannels 0 := by decide) :

          Build a minimal epsilon-predictor conv net: conv -> relu -> conv -> relu -> conv -> relu -> conv.

          This stays compact enough for the eager CUDA example while giving the CIFAR trainer more denoising capacity than a bare two-layer network.

          Instances For
            def TorchLean.nn.models.epsResidualConvNet {d : } (cfg : EpsConvNetConfig d) (h_batch : cfg.batch 0 := by decide) (h_dataC : cfg.dataChannels 0 := by decide) (h_inC : cfg.dataChannels + 1 0 := by decide) (h_hiddenC : cfg.hiddenChannels 0 := by decide) :

            Build a stronger same-resolution residual epsilon predictor.

            Architecture:

            stem conv -> relu -> residual block -> relu -> residual block -> relu -> output conv

            Each residual block has shape hiddenC×H×W -> hiddenC×H×W and computes $x+\operatorname{conv}(\operatorname{relu}(\operatorname{conv}(x)))$. This compact residual denoiser omits U-Net downsampling, upsampling, and multi-scale skip concatenation. It is still a useful compact architecture because residual paths make the denoising problem much easier than a plain conv chain while staying within the eager CUDA memory envelope used by examples.

            Instances For

              Map a tensor from $[0,1]$ into the standard diffusion training range $[-1,1]$.

              Instances For

                Deterministic Gaussian epsilon tensor for an arbitrary diffusion shape.

                The (seed, step) pair is turned into the runtime RNG key, so examples and artifact generation can reproduce the same noising path without ambient randomness.

                Instances For
                  def TorchLean.diffusion.linearBeta (T : ) (betaStart betaEnd : Float) (t : ) :

                  linear beta schedule value at timestep t.

                  Instances For
                    def TorchLean.diffusion.alphaBarsLinear (T : ) (betaStart betaEnd : Float) :

                    Compute the cumulative products $\bar\alpha_t=\prod_{s\le t}(1-\beta_s)$ for a linear beta schedule.

                    These values connect clean data $x_0$, noised data $x_t$, and the epsilon target used by DDPM-style training.

                    Instances For
                      def TorchLean.diffusion.appendTimeChannel (leading : Spec.Shape) {d c : } (spatial : Vector d) (x : Spec.Tensor Float (leading.concat (Spec.Shape.ofList (c :: spatial.toList)))) (tNorm : Float) :
                      Spec.Tensor Float (leading.concat (Spec.Shape.ofList ((c + 1) :: spatial.toList)))

                      Append a constant time channel after arbitrary leading axes.

                      The input layout is (leading..., channels, spatial...). The result preserves every leading and spatial axis and changes only the channel count from c to c + 1.

                      Instances For
                        def TorchLean.diffusion.noisedSampleFromEps (leading : Spec.Shape) {d c : } (spatial : Vector d) (alphaBars : Array Float) (T : ) (x0 eps : Spec.Tensor Float (leading.concat (Spec.Shape.ofList (c :: spatial.toList)))) (step : ) :
                        Sample.Supervised Float (leading.concat (Spec.Shape.ofList ((c + 1) :: spatial.toList))) (leading.concat (Spec.Shape.ofList (c :: spatial.toList)))

                        Build an epsilon-prediction training sample from explicit noise.

                        The caller supplies eps, usually from the runtime RNG. Keeping randomness outside this helper makes the transformation reusable:

                        $x_t=\sqrt{\bar{\alpha}_t}\,x_0+\sqrt{1-\bar{\alpha}_t}\,\varepsilon$, with target $\varepsilon$.

                        Instances For
                          def TorchLean.diffusion.noisedSample (leading : Spec.Shape) {d c : } (spatial : Vector d) (alphaBars : Array Float) (T : ) (x0 : Spec.Tensor Float (leading.concat (Spec.Shape.ofList (c :: spatial.toList)))) (seed step : ) :
                          Sample.Supervised Float (leading.concat (Spec.Shape.ofList ((c + 1) :: spatial.toList))) (leading.concat (Spec.Shape.ofList (c :: spatial.toList)))

                          Build a deterministic epsilon-prediction training sample.

                          This is the common DDPM training step used by examples: draw reproducible Gaussian noise from (seed, step), corrupt $x_0$, and use that same noise as the target.

                          Instances For
                            def TorchLean.diffusion.ddimPrev {s : Spec.Shape} (abPrev ab : Float) (x_t epsHat : Spec.Tensor Float s) :

                            One deterministic DDIM reverse update ($\eta=0$).

                            Given $x_t$, predicted epsilon, and adjacent schedule values, this estimates $x_0$ and remixes it to the previous timestep.

                            We clamp the intermediate $x_0$ estimate to the training image range $[-1,1]$. This is the standard "clipped denoised" stabilizer used by many DDPM/DDIM samplers: without it, a compact model can drive one color channel far outside the data range and the final PPM exporter merely clips the damage into saturated color blobs.

                            Instances For

                              Write the first image in an RGB NCHW batch as an ASCII PPM.

                              This dependency-free writer emits portable image artifacts for examples and rendered diagnostics.

                              Instances For