# 🧠 Neighbor-Aware Localized Concept Erasure in Text-to-Image Diffusion Models

## 📝 Abstract

Concept erasure in text-to-image diffusion models seeks to remove undesired concepts while preserving overall generative capability. Localized erasure methods aim to restrict edits to the spatial region occupied by the target concept. However, we observe that suppressing a concept can unintentionally weaken semantically related neighbor concepts, reducing fidelity in fine-grained domains. We propose Neighbor-Aware Localized Concept Erasure (NLCE), a training-free framework designed to better preserve neighboring concepts while removing target concepts. It operates in three stages: (1) a spectrally-weighted embedding modulation that attenuates target concept directions while stabilizing neighbor concept representations, (2) an attention-guided spatial gate that identifies regions exhibiting residual concept activation, and (3) a spatially-gated hard erasure that eliminates remaining traces only where necessary. This neighbor-aware pipeline enables localized concept removal while maintaining the surrounding concept neighborhood structure. Experiments on fine-grained datasets (Oxford Flowers, Stanford Dogs) show that our method effectively removes target concepts while better preserving closely related categories. Additional results on celebrity identity, explicit content and artistic style demonstrate robustness and generalization to broader erasure scenarios.

---

## ✨ Features

- Localized attention masking based on concept-sensitive token importance
- Semantic filtering using dual embeddings (CLIP + Qwen)
- Projection-based concept removal with spectral expansion
- Two-pass attention modification on specific layers
- Supports multi-concept ablation:
  - Dog breeds
  - Flowers
  - Celebrities
  - Nudity

---

## 📁 Project Structure

```
NLCE/
├── object_erasure.py                # Dog/Flower pipeline
├── celebrity_erasure.py             # Celebrity-focused pipeline
├── evaluation/                      
│   ├── artist/
│   │   └──...                       # Files used for artist evaluation 
│   └── ...
├── data/                            # Config CSVs for all categories
│   ├── config_flowers_camellia.csv
│   └── ...
└── other/
    └── neighborhood_set_file_constructor/
        └── neighborhood_set_file_constructor.ipynb 
```

---

## 🧬 Method Overview

1. **Concept & Neighbor Selection**  
   Uses CLIP & SentenceTransformer dual embedding to retrieve semantically related concepts.

2. **Projection Matrix Construction**  
   Builds `P_c` to suppress the target and reinforce unrelated concepts via spectral weighting.

3. **Attention Intervention**  
   Replaces cross-attention forward pass using `WrappedProcessor` and `double_forward()` in `down_blocks.2.attentions.1`.

---

## 📦 Requirements

```bash
pip install -r requirements.txt
```

---

## 🚀 Usage

### Dog and Flower Concept Erasure

```bash
python object_erasure.py \
  --prompt_file data/demo.csv \
  --alpha 1.0 \
  --beta 1.0 \
  --demo_mode
```

> **Note:** When `demo_mode` is enabled, the system bypasses the neighbor-retrieval procedure and instead loads the predefined neighbor sets stored in `multi_neighbors.json`. This mode is to be used when the index file is not available (see **Notes** section).


### Celebrity Concept Erasure

```bash
python celebrity_erasure.py \
  --output_dir results_celebrity/outputs\
  --prompt_files celeb_configs \
  --alpha 1.0 \
  --beta 0.9
```


> **Note:** Make sure all config `.csv` files are inside one folder (`erasure_configs/`, `celeb_configs/`, etc.)

---


### I2P Dataset Concept Erasure

We provide a separate script, `I2P_generation.py`, to run NLCE-style concept erasure on the **I2P** dataset. This script is designed to work with a minimal interface – **do not pass any additional arguments beyond the ones shown below.**



For I2P, use:

```bash
python I2P_generation.py \
  --erasure_beta 1.0 \
  --preserve_gamma 0.5 \
  --thr 14

```
---


### Artist Erasure

```bash
python artist_erasure.py \
  --prompt_file data/big_artist_prompts.csv \
  --alpha 1.0 \
  --beta 1.0 \
```

---

### 📌 Notes

- The embedding index `dual_embedding_index_with_scores.pt` (~19GB) could not be included in the zip file due to space constraints. 
- A set of generated concept neighborhoods is given in `multi_neighbors.json` that can be used to generate images of Oxford Flowers.


