.github/README.md

🛡️ 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 (ICLR 2026).

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

Install

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

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:

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

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.

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

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

Citation

@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}
}