src/NeuralNetwork/flux2/__init__.py

"""Flux2 transformer model components for LightDiffusion-Next.

This module contains the Flux2 model architecture including:
- DoubleStreamBlock: Joint image-text transformer blocks
- SingleStreamBlock: Merged stream transformer blocks  
- Flux2: Main model class with forward pass

Ported from ComfyUI's Flux2 Klein implementation.
"""

from src.NeuralNetwork.flux2.model import Flux2, Flux2Params
from src.NeuralNetwork.flux2.layers import (
    DoubleStreamBlock,
    SingleStreamBlock,
    LastLayer,
    MLPEmbedder,
    EmbedND,
    QKNorm,
)

__all__ = [
    "Flux2",
    "Flux2Params",
    "DoubleStreamBlock",
    "SingleStreamBlock",
    "LastLayer",
    "MLPEmbedder",
    "EmbedND",
    "QKNorm",
]