# 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

```bash
uv add lpips-nnx
```

## Usage

```python
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]**:

```python
# 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:

```python
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.](https://github.com/richzhang/PerceptualSimilarity) (BSD-2-Clause)
- VGG16 weights from [timm](https://huggingface.co/timm/vgg16.tv_in1k) (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)
