Artifex: Generative Modeling Research Library

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 artifex

# With GPU support (CUDA 12.0+)
pip install artifex[cuda]

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

Train a Diffusion Model on Fashion-MNIST

import jax
import jax.numpy as jnp
import optax
from datarax import from_source
from datarax.core.config import ElementOperatorConfig
from datarax.dag.nodes import OperatorNode
from datarax.operators import ElementOperator
from datarax.sources import TfdsDataSourceConfig, TFDSSource
from flax import nnx

from artifex.generative_models.core.configuration import (
    DDPMConfig, NoiseScheduleConfig, UNetBackboneConfig,
)
from artifex.generative_models.models.diffusion import DDPMModel
from artifex.generative_models.training.trainers import (
    DiffusionTrainer, DiffusionTrainingConfig,
)

# 1. Load Fashion-MNIST with datarax
def normalize(element, _key):
    image = element.data["image"].astype(jnp.float32) / 127.5 - 1.0
    return element.replace(data={**element.data, "image": image})

source = TFDSSource(
    TfdsDataSourceConfig(name="fashion_mnist", split="train", shuffle=True),
    rngs=nnx.Rngs(0),
)
normalize_op = ElementOperator(ElementOperatorConfig(stochastic=False), fn=normalize, rngs=nnx.Rngs(1))
pipeline = from_source(source, batch_size=64) >> OperatorNode(normalize_op)

# 2. Create DDPM model
config = DDPMConfig(
    name="fashion_ddpm",
    input_shape=(28, 28, 1),
    backbone=UNetBackboneConfig(
        name="unet", in_channels=1, out_channels=1,
        hidden_dims=(32, 64, 128), channel_mult=(1, 2, 4), activation="silu",
    ),
    noise_schedule=NoiseScheduleConfig(
        name="cosine", schedule_type="cosine", num_timesteps=1000,
    ),
)
model = DDPMModel(config, rngs=nnx.Rngs(42))
optimizer = nnx.Optimizer(model, optax.adamw(1e-4), wrt=nnx.Param)

# 3. Configure trainer with SOTA techniques
trainer = DiffusionTrainer(
    noise_schedule=model.noise_schedule,
    config=DiffusionTrainingConfig(
        loss_weighting="min_snr", snr_gamma=5.0, ema_decay=0.9999,
    ),
)
jit_train_step = nnx.jit(trainer.train_step)

# 4. Training loop
rng = jax.random.PRNGKey(0)
for batch in pipeline:
    rng, step_rng = jax.random.split(rng)
    _, metrics = jit_train_step(model, optimizer, {"image": batch["image"]}, step_rng)
    trainer.update_ema(model)

# 5. Generate samples
samples = model.sample(n_samples_or_shape=16, steps=100)

See the Quickstart Guide for complete examples including VAE training, visualization, and more.

Next Steps

• Core Concepts

  ---

  Understand generative modeling concepts and Artifex architecture

  Learn

• First Model Tutorial

  ---

  Build your first generative model from scratch

  Tutorial

• 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