Skip to content

Artifex: Generative Modeling Research Library¤

License: MIT Python 3.10+ JAX Flax Code style: ruff

A research-focused modular generative modeling library built on JAX/Flax NNX, providing implementations of state-of-the-art generative models with multi-modal support and scientific computing focus.

Why Artifex?¤

  • State-of-the-Art Models

    VAEs, GANs, Diffusion, Flows, EBMs, Autoregressive, and Geometric models with 2025 research compliance

    Model Gallery

  • Research-Focused

    Hardware-aware optimization, distributed training, mixed precision validated through 2150+ tests

    Getting Started

  • Multi-Modal Support

    Native support for images, text, audio, proteins, and multi-modal data

    Modalities Guide

  • Scalable Architecture

    From single GPU to multi-node distributed training with FSDP and tensor parallelism

    Scaling Guide

Quick Start¤

Installation¤

# CPU-only version
pip install avitai-artifex

# With GPU support (CUDA 12.0+)
pip install "avitai-artifex[cuda12]"

See the Installation Guide for detailed setup instructions including Docker and source installation.

Start with the VAE Quickstart¤

The primary onboarding path is the checked-in VAE quickstart. It uses TFDSEagerSource, VAEConfig, VAETrainer, and train_epoch_staged to train on MNIST while keeping the hot path in JAX.

from datarax.sources import TFDSEagerSource
from datarax.sources.tfds_source import TFDSEagerConfig
from artifex.generative_models.core.configuration import DecoderConfig, EncoderConfig, VAEConfig
from artifex.generative_models.models.vae import VAE
from artifex.generative_models.training import train_epoch_staged
from artifex.generative_models.training.trainers import VAETrainer, VAETrainingConfig

See the Quickstart Guide for the full walkthrough and refer to the checked-in executable pair under docs/getting-started/quickstart.py and docs/getting-started/quickstart.ipynb in the repository.

Next Steps¤

  • Core Concepts

    Understand generative modeling concepts and Artifex architecture

    Learn

  • Quickstart Guide

    Train your first generative model with Artifex

    Quickstart

  • Model Guides

    Deep dive into VAEs, GANs, Diffusion, Flows, and more

    Guides

  • Examples

    Working examples for all model types and use cases

    Examples

Model Types¤

Model Type Best For Guide
VAE Latent space exploration, representation learning VAE Guide
GAN High-quality image synthesis, style transfer GAN Guide
Diffusion State-of-the-art generation, inpainting Diffusion Guide
Flow Exact likelihood, density estimation Flow Guide
EBM Compositional generation, constraints EBM Guide
Autoregressive Text, sequential data AR Guide
Geometric Proteins, molecules, 3D structures Examples

Architecture¤

See Architecture Overview for detailed system design, component structure, and extension points.

Citation¤

@software{artifex_2025,
  title = {Artifex: Generative Modeling Research Library},
  author = {Shafiei, Mahdi and contributors},
  year = {2025},
  url = {https://github.com/avitai/artifex},
  version = {0.1.0}
}

Contributing¤

We welcome contributions! See the Contributing Guide for guidelines.

GitHub | Issues | Discussions