# 🛡️ MemGuard: Memorization Guardrail for Diffusion Models

A lightweight guardrail that detects and mitigates training-data memorization
in Stable Diffusion outputs. Built on the metric from [Detecting and Mitigating Memorization in Diffusion Models through Anisotropy of the
Log-Probability](https://openreview.net/forum?id=HTPGy5ydAY) (ICLR 2026).

**▶ Live demo:** https://huggingface.co/spaces/asthanarohan/memguard


## Install

```bash
pip install -e .            # core library
pip install -e ".[api]"     # + FastAPI service
pip install -e ".[demo]"    # + Gradio demo
pip install -e ".[dev]"     # + tests / lint
```

## Library usage

```python
from memguard import MemorizationDetector, load_sd_pipeline

pipe = load_sd_pipeline()                 # SD v1-4 + DDIM scheduler
det = MemorizationDetector(threshold=0.9)
det.score(pipe, "The No Limits Business Woman Podcast")       # -> float in [0, 1]
det.is_memorized(pipe, "The No Limits Business Woman Podcast")# -> bool
det.check(pipe, "The No Limits Business Woman Podcast")
# -> {"prompt": ..., "score": ..., "memorized": ..., "threshold": 0.9}
```

### Mitigate a memorized generation:

```python
from memguard import GuardedDiffusionPipeline, MemorizationDetector, load_sd_pipeline

pipe = load_sd_pipeline()
guarded = GuardedDiffusionPipeline(pipe, MemorizationDetector(threshold=0.9))
result = guarded.generate("The No Limits Business Woman Podcast")
# -> {"image": <PIL.Image|None>, "score": float, "memorized": bool,
#     "mitigated": bool, "blocked": bool, "signal_before": float, "signal_after": float, ...}
```

## Service

```bash
pip install -e ".[api,diffusers]"
uvicorn app.api:app --reload
# POST /score  {"prompt": "..."}  ->  {"prompt", "score", "memorized", "threshold"}
# GET  /health
```

## Demo

Enter a prompt; the demo generates an image with SD v1-4 and shows it alongside
a memorization meter (flagged as memorized at/above the 0.9 threshold). Then, the memorization can be mitigated and a non-memorized image is generated.

```bash
pip install -e ".[demo,diffusers]"   # gradio + torch/diffusers (first run downloads SD v1-4)
python app/demo.py                   # local Gradio at http://127.0.0.1:7860
```

## Docker

```bash
docker build -t memguard .
docker run -p 8000:8000 memguard
```


## Citation

```bash
@inproceedings{
asthana2026detecting,
title={Detecting and Mitigating Memorization in Diffusion Models through Anisotropy of the Log-Probability},
author={Rohan Asthana and Vasileios Belagiannis},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=HTPGy5ydAY}
}
```