StableDiffusionPUPPI

Written by: Luke Johnson (sa21722)

Updated: 12/08/2024

Repository for the Bristol ML@L1 Hackathon project! :)

Getting started

Firstly, the ML stack used is PyTorch and 🤗 Hugging Face. Basic syntax and tutorials are available in the links.

The algorithm is a Stable Diffusion chain which uses the U-Net architecture as the Deep Learning Model at its core. See 🤗 Hugging Face diffusers library for further info.

The repo has a PyTorch version & Keras version. The PyTorch folder has all the notebooks working with a PyTorch pipeline.

An in-house implementation of U-Net is also available where model architecture can be changed in the following folders:

PyTorch Folder

• All the notebooks working with a PyTorch pipeline
• Note there are multiple PyTorch models with different architectures
• This PyTorch model is bug free (as far as I am aware)
• The other .py model files have bugs like unexplicitly defined variable names in forward pass, so may need some adjusting if you want to use them
• This folder stopped being updated after Week 6 of the slides

Keras Folder

• All the notebooks working with a Keras pipeline
• Note there are multiple Keras Models with different architectures
• These 2 Keras models: model1, model2 have great performance and a small size suitable for synthesis however their custom padding/ layers are not supported by hls4ml
• This model can be used for synthesising into firmware but performance is worse
• This folder stopped being updated after the end of the slides

★ Note: Keras 3.0 can now serve as API to TensorFlow, PyTorch & JAX if you prefer to use it.

★ This Powerpoint includes a summary of some of the main changes and results from the summer internship. These were the slides shown by Luke at the weekly Trigger Group meetings.

Setting up the env

Training Diffusion models is expensive; follow this set up on a node with access to a GPU, Bristol GPU04 is suitable (the code will automatically port training tasks to GPU if available).

1. Clone this repository
2. Create the conda environment with conda env create -f <env_name>.yaml
3. If using PyTorch use: env.yaml, if using Keras use: hls_env.yaml
4. The preferred method of interaction with this repository is via VSCode. Also download any extensions required for displaying jupyter notebooks as prompted.

Data & models

Root datasets can be found in: /cephfs/dice/projects/L1T/diffusion/datasets/

Already-processed numpy data can be found here: /storage/ek19824/stablediffusionpuppi/data (or can be loaded from .root and saved as .npy in current dir)

A sample trained huggingface model is here: /cephfs/dice/projects/L1T/diffusion/trained_models/trained_diffusor.pt

UNet2d model architecture and config can be changed in most of the <model>.py files

Large files

• It's a good idea to version trained models and their evaluation, this can be done with git Large File System (git LFS).

  • you can download git LFS with:

  curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh 
  mkdir -p ~/bin
  export PATH=~/bin:$PATH
  apt-get download git-lfs
  dpkg-deb -x git-lfs_* ~/bin

  • and add it your $PATH with:

  export PATH="$HOME/bin/usr/local/git-lfs/bin:$PATH"

• init git lfs in the repo with: git lfs install

Then you can use the standard git commands.

RECIPE

Interfacing with models (training, inference, jet clustering) is to be done via the jupyter notebooks provided. In these notebooks there are various abstractions defined in standalone python modules (some of which can also be run from the CLI as needed).

Additions should be pushed to a new branch followed by a Merge Request to master.

1. Training Models

• Use the training.ipynb notebook to train new subclassing API models and UNet2d

• The notebooks in Keras folder have a list of variables in the first cell you might need to change to work with your setup

• You can also adjust the saturation energy, the image dimensions and the number of training epochs etc.

• There is a separate training_hlsmodel.ipynb notebook designed to work with the functional API and different architecture in the model_for_hls.py model

• Note: training.ipynb notebook saves ONLY the weights to a .h5 file whereas training_hlsmodel.ipynb saves the model architecture AND weights to a .h5 file

2. Inference

• Use the inference.ipynb notebook to test the subclassing API models' ability to denoise an image after the training process is complete

• The notebooks in Keras folder have a list of variables in the first cell you might need to change to work with your setup

• inference.ipynb imports a model from its python file and applies the weights from an HDF .h5 file. This notebook should be used to run inference on subclassing API models and UNet2d

• inference_hls.ipynb imports model architecture AND weights from a .h5 file. This notebook is also designed to work with the functional API keras model and compiled HLS Model

3. Jet Finding

The jet_clustering.ipynb notebook can be used to:

• Run a jet finding algorithm to identify the local maxima in clean and denoised images

• Find the matching jets

• Calculate the Jet $P_t$ resolution distribution for $HH$ → $4b + 200PU$ vs. denoised events

• Identify the leading jet in each image

• Plot how the Jet Rate decreases as you increase the $P_t$ cut

The jet finding algorithm can be altered in this file

Note: This notebook currently only works with the subclassing API models and UNet2d

4. HLS Conversion

Synthesis of the model for FPGA deployment, including firmware development using HLS4ML with hls_converter.py

• Loads a trained Keras model (which should be the one labelled model_for_hls.py)

• Sets up the environment for VivadoHLS 2019.2 (hls4ml supports Vivado versions 2018.2 to 2020.1 and Vitis HLS 2020.2 to 2022.2 experimentally)

• Creates a shell script to export necessary paths and environment variables for Vivado

• Makes the shell script executable and sources it to update the environment

• Converts a Keras model to an HLS model using hls4ml and compiles it

• Configures the HLS model with precision settings and optimization strategies

• Generates and saves layer-wise comparison plots between the Keras model and the HLS model

• Builds the HLS firmware by running C simulation, synthesis, and co-simulation steps

• HLS config saved to a JSON file and final firmware is generated

• Creates reports which contain info about performance, resources, latency, instance & loop details etc...

• Check output messages in hls_converter.log file

NOTE: Run the hls_converter.py script executing the shell script in the terminal with: ./run_hls_converter.sh

Tips & Ideas For Next Student

1. The numerical profiling function from hls4ml can be used to compare weights and activations between the Keras and HLS models in the hls_converter.py script. The X= argument should take in the inputs to the model which in our case is X = [noisy_images, pos_encoding, pos_encoding_bottleneck] but the positional encoding in the bottleneck has half the spatial dimensions as the positional encoding in the 1st block so the X = argument returns an error complaining of inhomogeneous input shapes.

   You will need to change the source code for the hls4ml profiling file. Look for the the numerical() function that calls another function: activations_hls_model(). This performs a trace with : _, trace = model.trace(np.ascontiguousarray(X)). The error comes from the ascontiguous numpy function being unable to broadcast arrays of different shapes together. Change the line to _, trace = model.trace(X).

2. Try implement custom padding (circular on sides and constant, 0, on top & bottom) in the keras model: model_for_hls.py without using custom layers or unsupported TFOpLambda layers/ custom tf operations (unsupported by hls4ml)

3. Play around with increasing model complexity while satisfying FPGA resources constraints

4. Adjust the hls config to reduce latency