Introduce Diffusion Models – Part 2 (Extended Models & Schedulers)

[ooples]

User Story

As a machine learning engineer, I want to leverage a wider variety of diffusion schedulers and advanced diffusion models (like Latent Diffusion) within the AiDotNet library, so that I can experiment with cutting-edge techniques for generative AI and improve the quality and speed of image generation.

Architectural Foundation (Assumed from #261 & #263)

This plan assumes the foundational interfaces and models from Part 1 exist, specifically: IDiffusionModel<T>, IStepScheduler<T>, DDPMModel<T>, and DDIMScheduler<T>.

---

Phase 1: Advanced Scheduler Implementation

Goal: Implement a suite of popular and powerful diffusion schedulers.

AC 1.1: PNDMScheduler<T> (3 points)

• Create file: src/Diffusion/Schedulers/PNDMScheduler.cs.
• Implement PNDMScheduler<T> to conform to the IStepScheduler<T> interface.
• The Step method must correctly implement the Pseudo Numerical Methods for Diffusion Models algorithm, which involves using the outputs from the previous 4 steps to calculate the next one.
• Add unit tests verifying the scheduler's output against a known Python implementation for the first 5 steps.

AC 1.2: DPMSolverMultistepScheduler<T> (5 points)

• Create file: src/Diffusion/Schedulers/DPMSolverMultistepScheduler.cs.
• Implement DPMSolverMultistepScheduler<T> to conform to IStepScheduler<T>.
• The Step method must implement the DPM-Solver++ algorithm, including the data-dependent correction term.
• Add unit tests that verify the single-step and multi-step outputs against a reference implementation.

AC 1.3: Additional Schedulers (3 points each)

• LMSDiscreteScheduler<T>: Implement the Linear Multi-Step scheduler.
• HeunDiscreteScheduler<T>: Implement the Heun's method second-order ODE solver.
• EulerDiscreteScheduler<T>: Implement the simple first-order Euler method.
• EulerAncestralDiscreteScheduler<T>: Implement the Euler method with added noise at each step.
• Each scheduler must be in its own file, implement IStepScheduler<T>, and have corresponding unit tests.

---

Phase 2: U-Net and VAE Implementation

Goal: Build the core neural network architectures required for latent diffusion.

AC 2.1: U-Net Sub-Components (8 points)

Requirement: Create the generic building blocks of the U-Net.

• Create a new folder: src/NeuralNetworks/Architectures/UNet/
• ResNetBlock<T>:
  • Create ResNetBlock.cs.
  • Implement a residual block with two generic Convolution layers, two GroupNormalization layers, and a SiLU activation. It must handle a skip connection.
• AttentionBlock<T>:
  • Create AttentionBlock.cs.
  • Implement a standard self-attention mechanism using generic Linear layers for Query, Key, and Value projections, followed by a Softmax and matrix multiplications.
• DownsampleBlock<T> & UpsampleBlock<T>:
  • Create these files. DownsampleBlock will be a Convolution layer with stride > 1. UpsampleBlock will be a TransposedConvolution layer.

AC 2.2: U-Net Model (8 points)

Requirement: Assemble the U-Net from its building blocks.

• Create UNet.cs in the same folder.
• The UNet<T> class will not be a traditional model but an architecture definition.
• The constructor will define the network layers, creating instances of the blocks from AC 2.1 in a sequence (e.g., a list of down-blocks, a mid-block, a list of up-blocks).
• Implement the Forward(Tensor<T> input, Tensor<T> timestepEmbedding) method:
  • It must pass the input through the downsampling blocks, storing the outputs for skip connections.
  • It must pass the result through the middle block.
  • It must pass the result through the upsampling blocks, concatenating the skip connection tensors at each stage.

AC 2.3: VAE Model (8 points)

Requirement: Build the Variational Autoencoder to move between pixel and latent space.

• Create a new folder: src/NeuralNetworks/Architectures/VAE/
• VAEEncoder<T>:
  • Create VAEEncoder.cs. Implement a series of Convolution layers that progressively reduce the spatial dimensions of an input Tensor<T>.
• VAEDecoder<T>:
  • Create VAEDecoder.cs. Implement a series of TransposedConvolution layers that upsample a latent Tensor<T> back to image dimensions.
• VAE<T>:
  • Create VAE.cs. It will contain an instance of the encoder and decoder.
  • Implement an Encode(Tensor<T> image) method that returns a latent tensor.
  • Implement a Decode(Tensor<T> latent) method that returns an image tensor.

---

Phase 3: Latent Diffusion Pipeline

Goal: Combine all the components into a functioning latent diffusion model.

AC 3.1: LatentDiffusionModel<T> (5 points)

Requirement: Create the final pipeline class that orchestrates the full generative process.

• Create file: src/Models/Generative/Diffusion/LatentDiffusionModel.cs.
• The class LatentDiffusionModel<T> must implement IDiffusionModel<T>.
• The constructor must accept a UNet<T>, a VAE<T>, and an IStepScheduler<T>.
• Implement the Generate method:
  • Step 1: Generate an initial random tensor (noise) with the target latent space dimensions.
  • Step 2: Loop from t=0 to num_inference_steps.
  • Step 3 (U-Net): In each loop, call the UNet's Forward method to predict the noise for the current latent tensor.
  • Step 4 (Scheduler): Call the scheduler's Step method, passing in the U-Net's output and the current latent tensor, to get the denoised latent tensor for the next step.
  • Step 5 (VAE): After the loop finishes, call the VAE's Decode method on the final latent tensor.
  • Step 6: Return the resulting image tensor from the decoder.

Definition of Done

• All checklist items are complete.
• All new schedulers and models are implemented and have unit tests against reference outputs.
• The LatentDiffusionModel is covered by an integration test that ensures all three main components (U-Net, VAE, Scheduler) are called in the correct order.
• All new code meets the project's >= 90% test coverage requirement.
• All new code is fully generic and adheres to all project architecture rules.

---

⚠️ CRITICAL ARCHITECTURAL REQUIREMENTS

Before implementing this user story, you MUST review:

• 📋 Full Requirements: .github/USER_STORY_ARCHITECTURAL_REQUIREMENTS.md
• 📐 Project Rules: .github/PROJECT_RULES.md

Mandatory Implementation Checklist

1. INumericOperations Usage (CRITICAL)

• Include protected static readonly INumericOperations<T> NumOps = MathHelper.GetNumericOperations<T>(); in base class
• NEVER hardcode double, float, or specific numeric types - use generic T
• NEVER use default(T) - use NumOps.Zero instead
• Use NumOps.Zero, NumOps.One, NumOps.FromDouble() for values
• Use NumOps.Add(), NumOps.Multiply(), etc. for arithmetic
• Use NumOps.LessThan(), NumOps.GreaterThan(), etc. for comparisons

2. Inheritance Pattern (REQUIRED)

• Create I{FeatureName}.cs in src/Interfaces/ (root level, NOT subfolders)
• Create {FeatureName}Base.cs in src/{FeatureArea}/ inheriting from interface
• Create concrete classes inheriting from Base class (NOT directly from interface)

3. PredictionModelBuilder Integration (REQUIRED)

• Add private field: private I{FeatureName}<T>? _{featureName}; to PredictionModelBuilder.cs
• Add Configure method taking ONLY interface (no parameters):

  public IPredictionModelBuilder<T, TInput, TOutput> Configure{FeatureName}(I{FeatureName}<T> {featureName})
  {
      _{featureName} = {featureName};
      return this;
  }

• Use feature in Build() with default: var {featureName} = _{featureName} ?? new Default{FeatureName}<T>();
• Verify feature is ACTUALLY USED in execution flow

4. Beginner-Friendly Defaults (REQUIRED)

• Constructor parameters with defaults from research/industry standards
• Document WHY each default was chosen (cite papers/standards)
• Validate parameters and throw ArgumentException for invalid values

5. Property Initialization (CRITICAL)

• NEVER use default! operator
• String properties: = string.Empty;
• Collections: = new List<T>(); or = new Vector<T>(0);
• Numeric properties: appropriate default or NumOps.Zero

6. Class Organization (REQUIRED)

• One class/enum/interface per file
• ALL interfaces in src/Interfaces/ (root level)
• Namespace mirrors folder structure (e.g., src/Regularization/ → namespace AiDotNet.Regularization)

7. Documentation (REQUIRED)

• XML documentation for all public members
• <b>For Beginners:</b> sections with analogies and examples
• Document all <param>, <returns>, <exception> tags
• Explain default value choices

8. Testing (REQUIRED)

• Minimum 80% code coverage
• Test with multiple numeric types (double, float)
• Test default values are applied correctly
• Test edge cases and exceptions
• Integration tests for PredictionModelBuilder usage

---

⚠️ Failure to follow these requirements will result in repeated implementation mistakes and PR rejections.

See full details: .github/USER_STORY_ARCHITECTURAL_REQUIREMENTS.md