DALL-E Mini Running in the Browser (work in progress)

Notes:

• Working tflite conversion in this notebook.
  • Note that the encoder and decoder need to be converted separately for some reason. More info on this bug.
  • But these models currently require TF Select operators due to bitwise operations that aren't supported in the default tflite opset. This issue tracks the potential inclusion of TF Select ops in tfjs-tflite, so models like this can be used on the web.
  • The encoder requires the TF Select Erf op. The decoder requires these TF Select ops: Bitcast, BitwiseOr, BitwiseXor, ConcatV2, Erf, LeftShift, RightShift, ScatterNd, SelectV2, StridedSlice
  • Here are the two models:
    • dalle_bart_encoder_flax_model
    • dalle_bart_decoder_flax_model
• ONNX conversion is blocked on support for Bitcast, BitwiseOr and BitwiseXor: https://github.com/onnx/tensorflow-onnx/issues/1985 The need for the BitwiseAnd op can be avoided by pulling jax.random.PRNGKey outside of the encode-decode function, but I'm not sure where the other ops are coming from - netron doesn't seem to "expand" certain types of nodes, so I think they're hidden within some other nodes.
• TF.js conversion is blocked on lack of support for these ops: BitwiseXor, Bitcast, LeftShift, BitwiseOr, RightShift (relevant issue)