Converting Keras CV Stable Diffusion models to tflite

environments

TensorFlow: 2.12.1 keras: 2.12.0 keras-cv: 0.4.2

Note that different version of TensorFlow and Keras could generate incompatible tflite files.

Conversion scripts

fp32

Tested with TensorFlow master branch (supposedly, using tf-nightly should be fine)

0. update group normalization, if you run into related problems, keras-team/keras-cv#1035
1. text encoder and decoder: Converting the text_encoder and decoder is trivial.
2. diffusion model: Converting the diffusion model needs extra effort. The model weights of the diffusion models is about 3.4 GiB, which is much larger than file size limit (2 GiB) of flatbuffer, the format TFLite used for its models. Surely, it's possible to modify flatbuffer to use 64-bit offset so that we can overcome the 2 GiB limit, but that will result in imcompatible files.

full-integer quantization

It's also possible to generate Quantized int8 models with Post-Training Quantization (PTQ). I don't know how to generate representative datasets, but as a proof of concept, I wrote a script that uses only one sample input as the dataset :-)

dynamic range quantization

Conversion script

Testing/Verifying converted tflite models

Borrowing some code from Keras CV implementation, we can do end-to-end test of converted TFLite models

• fp32 modeles: notebook
• dynamic range quantized modeles: notebook

test image encoder model with inpainting

• dynamice range quantized models: notebook The 'man-on-skateboard-cropped.png' is from this tutorial

Converted models

I put tflite models I converted to HuggingFace