This guide provides step-by-step instructions to reproduce the experiments and results presented in our Flashps paper. Follow these steps to validate our major claims regarding performance improvements and image quality preservation.
To simplify reproducibility, we provide an off-the-shelf Docker image, jiangxiaoxiao/flashps which includes all the dependencies and configurations required to run the experiments. This eliminates the need for complex environment setup.
We also provide an AWS EC2 instance with A10 GPUs.
Please comment on us in the HotCRP with your public key to get the IP address and access the machine.
You can refer to these commands to build the Docker environment and set up the conda environment. First, use the image pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel as the base image, then configure the conda environment within this image.
# Create conda env
conda create -n flashps python=3.10
conda activate flashps
# Git clone this repo
git clone https://github.com/Sylvia-16/FlashPS.git
pip install -r requirements.txt
# Install our customized diffusers package
cd diffusers && pip install -e .We have pulled the image on the provided machine, as its size is nearly 100 GiB.
# We will add your public key to the machine. You can log in to the machine with your private key by replacing the <IP_ADDRESS> with the actual IP.
ssh ubuntu@<IP_ADDRESS>
# We have pulled the image on the provided machine. You can skip this. On your machine, you can pull the prebuilt Docker image with the following command.
docker pull jiangxiaoxiao/flashps:latest
# Clear the stopped container, if it exists
docker kill flashps-ae
docker rm flashps-ae
# Run the following command to spin up the container. This may take a few minutes.
docker run -d --name flashps-ae --runtime=nvidia --gpus all --shm-size=16g \
-e NVIDIA_VISIBLE_DEVICES=all -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
-e CONDA_DEFAULT_ENV="" \
-e CONDA_AUTO_ACTIVATE_BASE=false \
jiangxiaoxiao/flashps sleep infinity
docker exec -it flashps-ae zsh
# Activate the environment
conda activate flashps
cd /app/image-inpainting/scheduler/
# Ensure the repo up-to-date
git pull
# Spin up the server to test TeaCache and diffusers baseline. It may take two minutes to start the server.
# When the server successfully starts, it will print
# "INFO: Uvicorn running on http://0.0.0.0:8005 (Press CTRL+C to quit)"
# on the console.
bash run_server_ootd_no_cb.sh
# Send requests to the server. Note that the first 10 requests are for warm-up purposes.
# For each baseline, we send requests with different RPS.
# Send requests to evaluate the baseline TeaCache.
bash /app/image-inpainting/scheduler/test_ootd_teacache.sh
# Send requests to evaluate the baseline diffusers.
bash /app/image-inpainting/scheduler/test_ootd_diffusers.sh
# Remember to kill the server.
bash kill_server.sh
# Spin up the server to test FlashPS baseline. It may take two minutes to start the server.
# When the server successfully starts, it will print
# "INFO: Uvicorn running on http://0.0.0.0:8005 (Press CTRL+C to quit)"
# on the console.
bash run_server_ootd.sh
# Send requests to evaluate FlashPS.
bash /app/image-inpainting/scheduler/test_ootd_flashps.sh
# Remember to kill the server.
bash kill_server.sh
# Analyze and plot the results. The script will print out the path to the figure.
python scripts/parse_end2end.py
You may compare the figure with Figure 12 in the paper.
Our image quality evaluation ensures that performance optimizations do not compromise output quality. As generating all images takes hours, we have cached them for evaluation.
Evaluate Image Quality. This may take minutes.
cd /app/image-inpainting/
bash scripts/test_quality.sh
You may compare the printed results with those in Table 2 in the paper.
As FISEdit requires a customized environment to execute, we have preconfigured the environment to evaluate the performance of SD2 directly on the machine (not using a Docker container). We configured the FISEdit environment following this link: https://github.com/Hankpipi/diffusers-hetu. The hetu.exp mentioned later in the script is the inference engine used by FISEdit.
Please log out and re-log in to the machine for the following evaluation.
ssh ubuntu@<IP_ADDRESS>
# Then stay on the machine, without the need to use a Docker container.# Initialize the environment
source activate pytorch
# Go to the project directory
cd /home/ubuntu/image-inpainting/scheduler
# Spin up the server to evaluate FlashPS. It may take two minutes to start the server.
# When the server successfully starts, it will print
# "INFO: Uvicorn running on http://0.0.0.0:8005 (Press CTRL+C to quit)"
# on the console.
bash run_server_sd2_cb.sh
# Send requests to evaluate FlashPS.
bash scripts/test_cb_sd2.sh
# kill the server
bash scripts/kill_gpu_processes.sh
# Spin up the server to evaluate Diffusers. It may take two minutes to start the server.
# When the server successfully starts, it will print
# "INFO: Uvicorn running on http://0.0.0.0:8005 (Press CTRL+C to quit)"
# on the console.
bash run_server_sd2_no_cb.sh
# Send requests to evaluate Diffusers.
bash scripts/test_no_cb_sd2.sh
# kill the server
bash scripts/kill_gpu_processes.sh
# activate fisedit environment
conda activate fisedit
source ~/Hetu/hetu.exp
# Spin up the server to evaluate FisEdit. It may take two minutes to start the server.
# When the server successfully starts, it will print
# "INFO: Uvicorn running on http://0.0.0.0:8005 (Press CTRL+C to quit)"
# on the console.
bash run_server_fisedit_no_cb.sh
# Send requests to evaluate FisEdit.
bash scripts/test_fisedit_e2e.sh
# kill the server
bash scripts/kill_gpu_processes.sh
# analyze and plot the result. The script will print out the path to the figure.
python scripts/parse_end2end.py
You may compare the figure with Figure 12 in the paper.
In the same machine, which is used to evaluate SD2.
# Go to the trace directory
cd /home/ubuntu/plot_mask_ratio
# Plot the distribution
python3 plot_mask_ratio_2traces.pyYou may compare the plot /home/ubuntu/plot_mask_ratio/mask_ratio_2traces.pdf with the Figure 3 in the paper.