Feature/diffusion updates

[opooladz]

Add comprehensive image diffusion training capabilities to EasyDeL

This PR introduces full-featured image diffusion training support with multiple architectures and trainers:

New Model Architectures

DiT (Diffusion Transformer)

• Clean implementation of diffusion transformers with patch-based image encoding
• Adaptive layer norm for timestep and class conditioning
• Supports multiple hidden sizes and depth configurations

DiT-MoE (Diffusion Transformer with Mixture of Experts)

• Based on DeepSeek V3 architecture (256 routed experts + 1 shared expert)
• Sigmoid scoring function with token-choice routing
• Configurable expert activation (top-k selection per token)
• Integrates with EasyDeL's BaseMoeModule for consistent MoE handling

UNet2D

• Full UNet architecture for Stable Diffusion training
• ResNet blocks, attention blocks, downsampling/upsampling
• Cross-attention support for text conditioning
• Time and class embedding layers

VAE (Variational Autoencoder)

• Encoder/decoder architecture for latent diffusion
• KL divergence regularization
• Seamless integration with Stable Diffusion pipeline

FLUX

• Modern diffusion architecture for high-quality generation
• Modulated attention and feedforward blocks

New Trainers

ImageDiffusionTrainer

• Trains DiT and DiT-MoE models on image datasets
• Rectified flow training (velocity prediction: v = data - noise)
• Min-SNR weighting (γ=5.0) for improved training stability
• Per-example RNG keys for reproducible noise generation
• Configurable timestep sampling and prediction types

StableDiffusionTrainer

• Text-to-image training with CLIP text encoder
• VAE latent space encoding
• UNet denoising in latent space
• Supports classifier-free guidance during training

Key Features

• Deterministic Training: Per-example RNG handling ensures reproducible results
• Flexible Prediction Types: Support for epsilon, velocity, and sample prediction
• Modern Training Techniques: Min-SNR weighting, rectified flow, loss scaling
• EasyDeL Integration: Uses EasyDeLState, Flax NNX modules, and standard trainer base classes
• Comprehensive Testing: Unit tests for all models and trainers
• Production Ready: Proper error handling, type hints, and documentation

Architecture Upgrade: DeepSeek V2 → V3

DiT-MoE now uses DeepSeek V3's improved MoE design:

• 256 routed experts (up from 160 in V2)
• Sigmoid scoring instead of softmax for better expert utilization
• Token-choice routing for efficient expert selection
• Dedicated shared expert for common features

Example Usage

import easydel as ed

# DiT-MoE with DeepSeek V3 architecture
model = ed.AutoEasyDeLModelForImageDiffusion.from_pretrained(
    "dit_moe",
    config=ed.DiTMoEConfig(
        image_size=32,
        patch_size=2,
        hidden_size=1152,
        num_hidden_layers=28,
        n_routed_experts=256,
        num_experts_per_tok=8,
        n_shared_experts=1,
        scoring_func="sigmoid",
    ),
)

trainer = ed.ImageDiffusionTrainer(
    model=model,
    arguments=ed.ImageDiffusionConfig(
        output_dir="./dit_moe_output",
        prediction_type="velocity",
        min_snr_gamma=5.0,
        num_train_timesteps=1000,
        learning_rate=1e-4,
    ),
    train_dataset=train_dataset,
)

trainer.train()

This brings EasyDeL's capabilities beyond LLMs into the image generation domain while maintaining the same high-quality training infrastructure.