# Evolutionary Caching to Accelerate Your Off-the-Shelf Diffusion Model

## Requirements

Follow these steps to set up the environment and install all necessary dependencies:

1. **Create a virtual environment**

   We use Conda with the provided `ecad.yml` to set up our environment, but you're free to use Python's built-in `venv` or any other environment manager you prefer.

2. **Create the Conda environment** using the provided environment file:

   ```bash
   conda env create -f ecad.yml
   ```

3. **Activate the environment**:

   ```bash
   conda activate ecad
   ```

4. **Install Python dependencies** via pip:

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

> Ensure that you are using Python 3.10 or later as specified in `ecad.yml`.

### Installing Local calflops Package

We have included a local copy of the `calflops` repository containing modifications from the original for an anonymous submission. Install it directly from the directory or archive:

```bash
# make sure your new env is activated
conda activate ecad

# replace this to the path to the directory if you have moved it
pip install ../calculate-flops.pytorch/
```

5. **Install the package in editable mode**:

   ```bash
   pip install -e .
   ```

   The last command installs the `editing-DiTs` directory as a package in 'edit' mode. If you run into issues with relative imports, rerun this command from the git root.

6. **Configure environment variables**:

   Modify the `.env` file to match your local setup. Update any paths accordingly. For example:

   ```bash
   # In .env
   PROJECT_ROOT="~/dev/editing-DiTs"  # SET THIS TO YOUR PROJECT ROOT, i.e., where this directory resides
   ```

   Replace the example path with the actual location of your project directory on your system.

## ECAD Optimization

ECAD optimization is achieved by running the `train_nsga2.py` script, which implements the NSGA‑II algorithm and is designed for parallel execution on SLURM clusters. Because this script leverages SLURM for parallel execution, running it directly on a non-cluster machine may be challenging and require code modification.

We **recommend** using `nohup` when running on a remote server to protect against SSH connectivity issues.

Below is an example command to optimize the PixArt-α model at 256×256 resolution:

```bash
nohup python editing_DiTs/genetic/train_nsga2.py \
    --image-generator PixArtAlphaImageGenerator \
    --name pixart_alpha_256_random_init \
    --num-cycles 500 \
    --all-populations-dir results/genetic/pixart_alpha/pixart_alpha_256_random_init \
    --all-benchmarks-dir results/benchmark/genetic/pixart_alpha/pixart_alpha_256_random_init \
    --batch-size 100 \
    --population-size 72 \
    --embedding-dir /path/to/precomputed/embeddings/for/below/prompts/image_reward_on_alpha_t5 \
    --benchmark-prompts prompts/image-reward-prompts.json \
    &> nohup_pixart_alpha_ecad_optimize.out & disown
```

To **resume** optimization from a previous checkpoint, replace the `--name` flag with:

```bash
--load-from results/genetic/pixart_alpha/pixart_alpha_256_random_init/gen_010/manager_config.json
```

For a full list of options and parameters, simply append the `--help` flag:

```bash
python editing_DiTs/genetic/train_nsga2.py --help
```

## ECAD Schedules

All caching schedules are stored as JSON files that model the binary tensor format detailed in the paper. The schedules referenced in the paper can be found under:

```
schedules/schedules_in_paper
```

When generating images, you will need to provide a path to a schedule file, or a whole directory of schedules for each of which image will be generated.

Note that resolution, guidance scale, custom weights, and custom pipelines are all manged through these JSON files. 

Some schedules used to create heuristics are also included in the `schedules` directory.

## Preparing Prompts

All prompt files used in the paper are available in the `prompts/` directory:

* `COCO_caption_prompts_30k.txt`
* `DrawBench200.txt`
* `ImageRewardPrompts.txt`
* `PartiPrompts.txt`
* `MJHQ-30K-prompts.json`
* `samples.txt` (provided by the PixArt-α authors, unmodified; may contain typos)

We include prompts to generate images included in qualitative figures of the paper in this directory too.

Before running inference, generate and save embeddings for your desired prompts using `generate_embeddings.py`. For example:

```bash
python editing_DiTs/benchmark/generate_embeddings.py \
    --image-generator PixArtAlphaImageGenerator \
    --file-path prompts/ImageRewardPrompts.txt \
    --output-dir embeddings/alpha_image_reward \
    --batch-size 32
```

This command processes the prompts in `prompts/ImageRewardPrompts.txt` with the specified image generator, saving embedding files to `embeddings/alpha_image_reward`. Adjust `--batch-size` as needed.

For detailed usage, run:

```bash
python editing_DiTs/benchmark/generate_embeddings.py --help
```

## Image Generation

Once embeddings are prepared, use `generate_images.py` to render images under each caching schedule. The script automatically creates a subdirectory per schedule in the output folder. Example command:

```bash
python editing_DiTs/benchmark/generate_images.py \
    --image-generator PixArtAlphaImageGenerator \
    --input-dir embeddings/alpha_image_reward \
    --schedule-dir schedules/schedules_in_paper \
    --output-dir outputs/images/alpha_image_reward \
    --num-images-per-prompt 10 \
    --batch-size 100 \
    --seed 0 \
    --seed-step 1
```

This will process each prompt embedding in `embeddings/alpha_image_reward`, apply every JSON schedule found in `schedules/optimized_schedule`, and save the results into `outputs/images/alpha_image_reward/<schedule-name>/`. Adjust `--num-images-per-prompt`, `--batch-size`, or other flags as needed.

For full parameter details, run:

```bash
python editing_DiTs/benchmark/generate_images.py --help
```

## Benchmarking

### Computing MACs

After generating images, you can compute static complexity metrics—such as Multiply–Accumulate operations (MACs) and FLOPs—for each caching schedule using the `compute_macs.py` script. This is useful for understanding the computational cost of different schedules without running full inference.

Example command:

```bash
python editing_DiTs/benchmark/compute_macs.py \
    --image-generator PixArtAlphaImageGenerator \
    --input-dir schedules/schedules_in_paper \
```

Here:

* `--input-dir` specifies a directory (or single JSON file) of schedules to analyze. The computed values are dumped into each JSON schedule that it checks. Open it to see the results.
* Use `--recompute-existing` to force re‑calculation of metrics for schedules that already have results.

For full usage information, run:

```bash
python editing_DiTs/benchmark/compute_macs.py --help
```

### Latency Computation

To measure the actual runtime latency of each caching schedule on your prompt embeddings, use the `compute_latency.py` script. This script runs inference repeatedly to warm up the GPU and averages over multiple batches.

Example command:

```bash
python editing_DiTs/benchmark/compute_latency.py \
    --image-generator PixArtAlphaImageGenerator \
    --embedding-dir embeddings/alpha_image_reward \
    --schedule-dir schedules/schedules_in_paper \
    --batch-size 100 \
    --warmup-steps 1 \
    --num-samples 5
```

* `--warmup-steps` and `--num-samples` should follow the values reported in the paper’s appendix to match the benchmarking protocol.
* You can also specify `--schedule-list` to measure a specific subset of schedule files.
* Use `--recompute-existing` to force re‑evaluation of latency for schedules with prior results.

For full usage details, run:

```bash
python editing_DiTs/benchmark/compute_latency.py --help
```

### Image Reward Scoring

After computing latency, evaluate the quality of your generated images using the ImageReward model. The `score_images.py` script processes images in a schedule-specific directory and attaches a score for each prompt.

Example command:

```bash
python editing_DiTs/benchmark/score_images.py \
    --benchmark-prompts prompts/ImageRewardPrompts.json \
    --image-dir outputs/images/alpha_image_reward/optimized_schedule \
    --output-subpath scores/image_reward_scores.json \
    --image-naming-mode image_reward \
    --file-mode json \
    --exactly-n-images 100
```

* `--benchmark-prompts`: JSON file of prompts (e.g., `prompts/MJHQ-30K-prompts.json`) or text file with one prompt per line when paired with `--file-mode txt`.
* `--image-dir`: Root directory containing images or subfolders of images per schedule.
* `--output-subpath`: Filename or nested path for the resulting score JSON (relative to a `scores/` base).
* `--image-naming-mode`: Choose from `image_reward`, `parti`, or `toca` to match how images are named. Details in the code's regex, but it is safe to always go with `parti`.
* `--file-mode`: `json`, `text`, `txt`, or `tsv` to match your prompt file format.
* `--exactly-n-images`: Only score directories containing exactly *n* images (prevents partial runs).
* Use `--delete-after` to remove images post-scoring and `--rescore-existing` to overwrite prior results.

For complete details, run:

```bash
python editing_DiTs/benchmark/score_images.py --help
```

### FID Computation

To assess the similarity between generated and reference images, compute the Fréchet Inception Distance (FID) using `compute_fid.py`. This measures the distributional distance between real and generated image features.

Example command:

```bash
python editing_DiTs/benchmark/compute_fid.py \
    --ref-dir data/MJHQ-30K \
    --gen-dir outputs/images/alpha_image_reward/optimized_schedule \

    --model-name inception_v3
```

* `--ref-dir`: Path to the reference MJHQ‑30K image directory. Note the directory structures must match between the reference and generated images directories.
* `--gen-dir`: Path to your generated images directory for a given schedule.
* `--model-name`: Inception variant to use (default: `inception_v3`).

For full usage details, run:

```bash
python editing_DiTs/benchmark/compute_fid.py --help
```

### CLIP Score Computation

Evaluate semantic alignment between your prompts and images using CLIP via `compute_clip.py`.

Example command:

```bash
python editing_DiTs/benchmark/compute_clip.py \
    --prompts-path prompts/MJHQ-30K-prompts.json \
    --image-dir outputs/images/alpha_image_reward/optimized_schedule \
    --file-mode json \
    --clip-model ViT-B/32 \
    --image-naming-mode image_reward
```

* `--prompts-path`: File containing the original prompts (JSON, text, or TSV).
* `--image-dir`: Directory with generated images to score.
* `--file-mode`: Format of the prompt file (`json`, `text`, `tsv`).
* `--clip-model`: Name of the CLIP model to use (e.g., `ViT-B/32`).
* `--image-naming-mode`: Pattern for matching image filenames (e.g., `image_reward`, `parti`, `toca`, `mjhq`, `coco`).

For full usage details, run:

```bash
python editing_DiTs/benchmark/compute_clip.py --help
```
