lpips-nnx 0.1.0

LPIPS (Learned Perceptual Image Patch Similarity) implementation in Flax NNX

LPIPS-NNX

LPIPS (Learned Perceptual Image Patch Similarity) implementation in Flax NNX, with weights that exactly match the official PyTorch implementation.

No PyTorch dependency at runtime - VGG16 weights are loaded from HuggingFace safetensors, LPIPS linear weights are bundled with the package.

Installation

uv add lpips-nnx

Usage

import flax.nnx as nnx
import jax.numpy as jnp
from lpips_nnx import LPIPS

# Initialize LPIPS (downloads VGG16 weights on first use, ~553MB cached)
lpips = LPIPS(rngs=nnx.Rngs(0))

# Create sample images (NHWC format, values in [-1, 1])
x = jnp.zeros((1, 224, 224, 3))
y = jnp.ones((1, 224, 224, 3)) * 0.5

# Compute perceptual distance
distance = lpips(x, y)
print(f"LPIPS distance: {distance.item():.4f}")

Input formats

Images should be in NHWC format (batch, height, width, channels) with values in [-1, 1]:

# If your images are in [0, 1] range, use normalize=True
distance = lpips(x_01, y_01, normalize=True)

JIT compilation

LPIPS is compatible with JAX's JIT compilation:

import jax

@jax.jit
def compute_lpips(x, y):
    return lpips(x, y)

distance = compute_lpips(x, y)

Dependencies

Runtime (lightweight):

• jax, jaxlib, flax
• safetensors - for loading VGG16 weights
• requests - for downloading weights

Development only:

• torch, torchvision, lpips - for comparison tests

Acknowledgments

• LPIPS weights from Richard Zhang et al. (BSD-2-Clause)
• VGG16 weights from timm (BSD-3-Clause)

References

Zhang, R., Isola, P., Efros, A. A., Shechtman, E., & Wang, O. (2018). The Unreasonable Effectiveness of Deep Features as a Perceptual Metric. CVPR.

License

MIT (see LICENSE for third-party attributions)