# AUTH-PROMPT BENCH: TOWARDS RELIABLE AND STABLE PROMPTING IN TEXT-TO-IMAGE GENERATION

This repository contains the code and models for the paper "Auth-Prompt Bench: Towards Reliable and Stable Prompting in Text-to-Image Generation" 

# Code Organization
1. `metrics/`: Contains evaluation metrics and related scripts.
2. `configs/`: Configuration files for different models and experiments.
3. `models/`: NoxEye Pipeline and related components (including CoMat).
4. `tools/`: Utility scripts for data processing and other tasks.
5. `data/`: Lexica and civitai datasets.
6. `README.md`: This file.
7. `requirements.txt`: List of dependencies.
8. `.gitignore`: Git ignore file.

# Getting Started
## Prerequisites
To get started, create a conda environment containing the required dependencies.
```bash
conda create -n auth-prompt-bench python=3.12
conda activate auth-prompt-bench
pip install -r requirements.txt
```

Create checkpoints directory and download the required models:
```bash
mkdir -p checkpoints
cd checkpoints
# For CoMat model
mkdir FastSAM
cd FastSAM
wget https://github.com/CASIA-IVA-Lab/FastSAM#model-checkpoints
python -m spacy download en_core_web_trf
cd ../..
```
## Training NoxEye
To train the NoxEye model, run the following command:
```bash
python tools/mistral_train.py --gpus=1 --nodes=1 --dataset_path data/diffusionDB/part-000001/train_prompt_caption_mainobj_v0-000001.json --name llm+obj --model_type custom --version 0
```
If you want to use CoMat, training with:
```bash
python scripts/comat_train.py  --gpus=4 --nodes=4 --config_path configs/comat/comat21.yaml
```

## Using Auth-Prompt Bench
To use the Auth-Prompt Bench for generating images based on prompts, run:
```bash
python -m metrics.attack --config-name=ATM/attack/mis_r_class_lora_v0+sd21 dataset_type=human data_path=data/lexica_V1_0.json end=20000 save_dir=./data/lexica/mis_r_class_lora_v0+sd21 other_kwargs.refiner_kwargs.refined_file="./data/lexica/mis_r_class_lora_v0+sd21/refined.csv"
```
Calculate energy and entropy:
```bash
python -m metrics.calc_energy_and_entropy --input_csv ./data/lexica/mis_r_class_lora_v0+sd21/refined.csv --output_csv ./data/lexica/mis_r_class_lora_v0+sd21/E_H.csv --model_name google/gemma-3-1b-it
```

## Using NoxEye for Inference
To use the NoxEye model for inference, run:
```python
from models.noxeye import NoxEyePipeline
model = NoxEyePipeline(refine_pretrain_model="mistralai/Mistral-7B-Instruct-v0.2+mllm")
model.to("cuda")
image = model.generate(["A photo of an astronaut riding a horse on mars", "A photo of a cat sitting on a chair"])
image[0].show()
```

