CGCE: Classifier-Guided Concept Erasure in Generative Models

Project Page | Paper

Viet Nguyen and Vishal M. Patel
Johns Hopkins University

Accepted at ECCV 2026.

Installation

Clone the repository and create a Python 3.12 environment:

git clone https://github.com/viettmab/CGCE.git
cd CGCE

conda create -n cgce python=3.12
conda activate cgce

Install PyTorch and the remaining dependencies:

pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 \
  --index-url https://download.pytorch.org/whl/cu126

python -m pip install "setuptools<81" wheel
python -m pip install --no-build-isolation -r requirements.txt

Some backbones require separately downloaded checkpoints, repository access, or model-license acceptance.

Train Concept Detectors

Train one detector by specifying a concept and text-encoder architecture:

bash scripts/train_concept.sh nudity sd14

Train several detector variants in sequence:

bash scripts/train_concept.sh nudity sd3_clip sd3_t5 flux_clip_pooled

Supported concepts include nudity, church, van_gogh, artist:<Name>, and custom. Supported architectures are:

• sd14
• sd3_clip
• sd3_clip_pooled
• sd3_t5
• flux_clip_pooled
• switti
• infinity_t5
• hunyuan_llava

For a custom concept, provide the dataset and concept text through environment variables:

DATASET_PATH=dataset_train/my_concept.json \
CONCEPT_TEXT="my concept" \
bash scripts/train_concept.sh custom sd14

BATCH_SIZE is per GPU, so a batch size of 8 on four GPUs produces a global batch size of 32.

Detector checkpoints are written to:

training-runs/<experiment-name>/concept_detector.pth

Image Generation

The unified image launcher supports Stable Diffusion 1.4, Stable Diffusion 3, FLUX, and Switti:

bash scripts/infer_cgce.sh sdv1
bash scripts/infer_cgce.sh sdv3
bash scripts/infer_cgce.sh flux
bash scripts/infer_cgce.sh switti

Select one or more prompt datasets with the space-separated DATASETS variable:

DATASETS="p4d sixcd" DEVICE=0 bash scripts/infer_cgce.sh sdv1

Supported dataset aliases are ring-a-bell, i2p, p4d, mma-diffusion, sixcd, and coco.

generate_cgce.py accepts detector checkpoints through one ordered argument:

--detector_checkpoints PATH [PATH ...]

The required checkpoint order is:

• SD 1.4 and Switti: CLIP
• SD3: T5, CLIP, pooled CLIP
• FLUX: T5, pooled CLIP

Generated images and logs are saved under:

results/<model>/<dataset>/<experiment>/
|-- all/
`-- logs.txt

Artist and object removal

The artist and object launchers use Stable Diffusion 1.4 only:

bash scripts/infer_cgce_artist.sh
bash scripts/infer_cgce_object.sh

Common inference settings can be overridden without editing the scripts:

DEVICE=1 STEP_SIZE=0.25 PERTURBATION_STRATEGY=raw_gradient \
  bash scripts/infer_cgce_artist.sh

Artist generation additionally creates erased/ and unerased/ directories for efficacy and preservation evaluation.

Video Generation

CogVideoX and HunyuanVideo use two dataset aliases:

• gen maps to dataset/nudity_gen_video.csv
• safesora-sexual maps to dataset/safe-sora.csv

Both datasets are generated by default:

bash scripts/infer_cgce_cogvideox.sh
bash scripts/infer_cgce_hunyuan.sh

Run only one dataset with:

DATASETS=gen bash scripts/infer_cgce_cogvideox.sh
DATASETS=safesora-sexual bash scripts/infer_cgce_hunyuan.sh

Both launchers explicitly enable concept removal. CogVideoX uses one T5 detector; HunyuanVideo uses one LLaVA/Llama detector and one pooled CLIP detector. Videos are saved as video_<row-id>.mp4.

Infinity Generation

Run Infinity concept removal with:

bash scripts/infer_cgce_infinity.sh

The launcher defaults to sixcd and supports ring-a-bell, i2p, p4d, mma-diffusion, and sixcd:

DATASETS="p4d sixcd" DEVICE=0 bash scripts/infer_cgce_infinity.sh

Infinity model and VAE checkpoints default to weights/infinity_2b_reg.pth and weights/infinity_vae_d32reg.pth.

Evaluation

Image nudity

Evaluate one or more CGCE save directories. Each directory must contain all/:

bash scripts/eval_nudity.sh \
  results/sdv1/p4d/concept_removal_importance_weighted_step_size1.0

For GPT-4o mini evaluation, set OPENAI_API_KEY and run:

OPENAI_API_KEY=<key> bash scripts/eval_nudity_gpt.sh

Artist removal

Artist evaluation requires a baseline save directory followed by one or more CGCE save directories. Each directory must contain matching erased/ and unerased/ images:

bash scripts/eval_artist.sh \
  results_baseline/sdv1/VanGogh \
  results/sdv1/VanGogh/concept_removal_importance_weighted_step_size0.5

This computes Van Gogh detection on erased/ and LPIPS preservation scores for both subsets.

Video nudity

bash scripts/eval_video.sh \
  results/cogvideox/gen/concept_removal_importance_weighted_step_size2.0

FID and CLIP score

bash scripts/fid_clip.sh <cgce-save-dir>

Set JOB=fid, JOB=clip, or JOB=both to select the metrics.

Citation

If you find CGCE useful, please cite:

@inproceedings{nguyen2026cgce,
  title={CGCE: Classifier-Guided Concept Erasure in Generative Models},
  author={Nguyen, Viet and Patel, Vishal M.},
  booktitle={European Conference on Computer Vision (ECCV)},
  year={2026}
}