| """SolPix-50M latent flow-transformer reference implementation.""" | |
| from .config import SolPixConfig | |
| from .autoencoder import AutoencoderDCSol | |
| from .model import SolPix, SolPixOutput, SolPixTransformer2D | |
| from .training import REPAHead, FlowMatchingConfig, flow_matching_loss, representation_alignment_loss | |
| __all__ = [ | |
| "SolPix", | |
| "SolPixTransformer2D", | |
| "AutoencoderDCSol", | |
| "SolPixConfig", | |
| "SolPixOutput", | |
| "FlowMatchingConfig", | |
| "REPAHead", | |
| "flow_matching_loss", | |
| "representation_alignment_loss", | |
| ] | |