Artifex
Modular Generative Modeling Library
A research-focused modular library for state-of-the-art generative models, built on JAX/Flax NNX for scientific reproducibility and experimentation.
Repository Coming Soon
This project is under active development and will be open-sourced soon
Overview
Artifex is a modular library for generative modeling research, providing implementations of state-of-the-art generative models with a focus on modularity, type safety, and scientific reproducibility. Built on JAX and Flax NNX, it emphasizes clean abstractions and extensible design for research experimentation.
The library is designed with a research-first philosophy — prioritizing modularity to make it easy to swap components and experiment, clarity through clean readable implementations, and extensibility to add new models and functionality. All configurations use frozen dataclasses with validation, and the entire codebase follows a protocol-based design with full type annotations.
Artifex supports a wide range of generative model families: VAEs, GANs, Diffusion Models, Normalizing Flows, Energy-Based Models, and Autoregressive Models, along with geometric models for point clouds, meshes, and protein structures. It provides unified modality support across images, text, audio, proteins, tabular data, time series, and multi-modal generation. The library includes modular losses, flexible sampling strategies, a domain extension system, and an evaluation framework for standardized benchmarking.
Key Features
VAE & Variational Models
VAE, β-VAE, VQ-VAE, and Conditional VAE with multiple encoder architectures (dense, CNN, ResNet) for structured latent representations.
GAN Models
DCGAN, WGAN, StyleGAN, CycleGAN, and PatchGAN architectures for high-quality generative modeling.
Diffusion Models
DDPM, DDIM, score-based models, DiT, and Latent Diffusion for high-fidelity generation with controllable sampling.
Normalizing Flows
RealNVP, Glow, MAF, IAF, and Neural Spline Flows for exact likelihood computation and invertible transformations.
Energy-Based Models
Langevin dynamics and MCMC sampling with BlackJAX integration for flexible probability modeling.
Multi-Modal Generation
Unified interface across images, text, audio, proteins, tabular data, and time series with cross-modal generation.
Use Cases
Protein structure generation and design for drug discovery
Multi-modal biological data generation across images, sequences, and structures
Small molecule generation for therapeutic compound discovery
Data augmentation for limited biological datasets
Geometric modeling for point clouds, meshes, and molecular structures
Transfer learning between biological domains
Anomaly detection in experimental data
Privacy-preserving tabular data generation
Installation
# Clone the repository
git clone https://github.com/avitai/artifex.git
cd artifex
# Run setup script (creates venv, installs dependencies, detects GPU)
./setup.sh
# Activate the environment
source ./activate.shQuick Start
import jax
from flax import nnx
from artifex.generative_models.models.vae import VAE
from artifex.generative_models.core.configuration import (
VAEConfig, EncoderConfig, DecoderConfig,
)
# Create configuration using frozen dataclasses
config = VAEConfig(
name="mnist_vae",
encoder=EncoderConfig(
name="encoder", input_shape=(28, 28, 1),
latent_dim=64, hidden_dims=(256, 128), activation="relu",
),
decoder=DecoderConfig(
name="decoder", output_shape=(28, 28, 1),
latent_dim=64, hidden_dims=(128, 256), activation="relu",
),
encoder_type="dense",
kl_weight=1.0,
)
# Initialize and run
model = VAE(config, rngs=nnx.Rngs(0))
batch = jax.random.normal(jax.random.key(0), (16, 28, 28, 1))
outputs = model(batch)
samples = model.sample(n_samples=16)Built With
Ready to Get Started?
Explore the documentation, try examples, or contribute to the project.