stable_diffusion_jax-to-tflite --- decode_latents.ipynb

stable_diffusion_jax-to-onnx.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "stable_diffusion_jax-to-tflite --- decode_latents.ipynb",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "premium"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/josephrocca/a88de2343afb946292e65ede5fafdcbd/stable_diffusion_jax-to-onnx.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# This notebook was created for this issue: https://github.com/tensorflow/tensorflow/issues/58125#issuecomment-1282886939"
      ],
      "metadata": {
        "id": "YS9ymUvxhbgJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install transformers==4.23.1 huggingface_hub==0.10.0 ftfy==6.1.1 flax==0.6.1 git+https://github.com/huggingface/[email protected] git+https://github.com/onnx/[email protected]\n",
        "!pip install --upgrade jax jaxlib\n",
        "#!pip install tensorflow==2.10.0rc0 # <-- Kernel crashes with this too.\n",
        "!pip install tf-nightly"
      ],
      "metadata": {
        "id": "0YHLndloz1U_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from huggingface_hub.hf_api import HfFolder\n",
        "HfFolder.save_token('h'+'f'+'_'+'AUxlCqSud'+'NTSgaWmE'+'jrUgRytG'+'JiBTLoYSD') # Don't worry! This key can be safely made public. It's just a read-only key for an \"empty\"/dummy Hugging Face account (temp email) that was SPECIFICALLY created to make it easier to access the Stable Diffusion model in Colab (less copy-pasting my token during many runtime resets). The `+` concatenation is just so it doesn't trigger any Github API key detection alarms, or whatever.\n",
        "import numpy as np\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "from PIL import Image\n",
        "import tensorflow as tf\n",
        "from jax.experimental import jax2tf\n",
        "from diffusers import FlaxStableDiffusionPipeline\n",
        "import tf2onnx"
      ],
      "metadata": {
        "id": "UFMtdmPeyxpi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# (It's safe to ignore the warning messages, everything is okay)\n",
        "pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\", revision=\"flax\", dtype=jnp.float32, safety_checker=None)"
      ],
      "metadata": {
        "id": "PPtraQX34Az7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vae_params = params[\"vae\"]\n",
        "\n",
        "def decode_latents(vae_params, latents):\n",
        "  latents = 1 / 0.18215 * latents\n",
        "  images = pipeline.vae.apply({\"params\": vae_params}, latents, method=pipeline.vae.decode).sample\n",
        "\n",
        "  images = (images / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1)\n",
        "  return images"
      ],
      "metadata": {
        "id": "m0nGZcx_ZAnU"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "decode_latents_closed_over_params = lambda latents: decode_latents(vae_params, latents)\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.experimental_from_jax([decode_latents_closed_over_params], [[('latents', np.ones([1, 4, 64, 64], dtype='float32'))]])\n",
        "\n",
        "tflite_model = converter.convert()\n",
        "with open('decode_latents.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)\n",
        "\n",
        "print(\"DONE!\")"
      ],
      "metadata": {
        "id": "8yd6TBH_TjRt"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}