flax_stable_diffusion_3.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/stable-diffusion-diffusers-colab/blob/main/flax_stable_diffusion_3.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Install - Step 1\n",
        "!pip install -q jax==0.3.25 jaxlib==0.3.25 flax==0.6.2 transformers accelerate omegaconf\n",
        "!pip install -q git+https://github.com/huggingface/diffusers\n",
        "\n",
        "!apt -y install -qq aria2\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M 'https://civitai.com/api/download/models/138176?type=Model&format=SafeTensor&size=pruned&fp=fp16' -d /content/models -o cyberrealistic_v33.safetensors\n",
        "# !aria2c --console-log-level=error -c -x 16 -s 16 -k 1M 'https://huggingface.co/lllyasviel/ControlNet-v1-1/resolve/main/control_v11p_sd15_canny.pth' -d /content/models/controlnet -o control_v11p_sd15_canny.pth\n",
        "# !aria2c --console-log-level=error -c -x 16 -s 16 -k 1M 'https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_canny_fp16.safetensors' -d /content/models/controlnet -o control_v11p_sd15_canny_fp16.safetensors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Install - Step 2\n",
        "import jax, random, gc, torch\n",
        "from flax.jax_utils import replicate\n",
        "from flax.training.common_utils import shard\n",
        "from diffusers import StableDiffusionPipeline\n",
        "\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "import jax.tools.colab_tpu\n",
        "jax.tools.colab_tpu.setup_tpu('tpu_driver_nightly')\n",
        "jax.device_count()\n",
        "\n",
        "def image_grid(imgs, rows, cols):\n",
        "    w,h = imgs[0].size\n",
        "    grid = Image.new('RGB', size=(cols*w, rows*h))\n",
        "    for i, img in enumerate(imgs): grid.paste(img, box=(i%cols*w, i//cols*h))\n",
        "    return grid"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Vanilla - Step 1\n",
        "pipeline = StableDiffusionPipeline.from_single_file('/content/models/cyberrealistic_v33.safetensors', torch_dtype=torch.float16, safety_checker=None)\n",
        "pipeline.save_pretrained('/content/models/torch/cyberrealistic', safe_serialization=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Vanilla - Step 2\n",
        "from diffusers import FlaxStableDiffusionPipeline\n",
        "\n",
        "pipeline, params = FlaxStableDiffusionPipeline.from_pretrained('/content/models/torch/cyberrealistic', from_pt=True, dtype=jax.numpy.bfloat16, safety_checker=None)\n",
        "params = replicate(params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Vanilla - Step 3\n",
        "gc.collect()\n",
        "\n",
        "prompt = \"cute duck\"\n",
        "negative_prompt = \"blurry\"\n",
        "\n",
        "num_samples = jax.device_count()\n",
        "\n",
        "prompt_n = num_samples * [prompt]\n",
        "prompt_ids = pipeline.prepare_inputs(prompt_n)\n",
        "prompt_ids = shard(prompt_ids)\n",
        "\n",
        "negative_prompt_n = num_samples * [negative_prompt]\n",
        "negative_prompt_ids = pipeline.prepare_inputs(negative_prompt_n)\n",
        "negative_prompt_ids = shard(negative_prompt_ids)\n",
        "\n",
        "real_seed = random.randint(0, 2147483647)\n",
        "prng_seed = jax.random.PRNGKey(real_seed)\n",
        "prng_seed = jax.random.split(prng_seed, jax.device_count())\n",
        "\n",
        "images = pipeline(prompt_ids, params, prng_seed, neg_prompt_ids=negative_prompt_ids, num_inference_steps=50, height=512, width=512, guidance_scale=7, jit=True).images\n",
        "images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))\n",
        "\n",
        "image = image_grid(images, 2, 4)\n",
        "image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title ControlNet #Canny - Step 1\n",
        "import jax, torch\n",
        "from diffusers import FlaxStableDiffusionControlNetPipeline, ControlNetModel\n",
        "# controlnet = ControlNetModel.from_single_file('/content/models/controlnet/control_v11p_sd15_canny.pth', dtype=torch.float16) # invalid load key, '<'.\n",
        "controlnet = ControlNetModel.from_single_file('https://huggingface.co/lllyasviel/ControlNet-v1-1/blob/main/control_v11p_sd15_canny.pth', dtype=torch.float16)\n",
        "pipeline, params = FlaxStableDiffusionControlNetPipeline.from_pretrained('/content/models/torch/cyberrealistic', controlnet=controlnet, dtype=jax.numpy.bfloat16, from_pt=True, safety_checker=None)\n",
        "params = replicate(params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title ControlNet #Canny - Step 2\n",
        "from diffusers.utils import load_image\n",
        "image = load_image(\"https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png\")\n",
        "image_raw = image\n",
        "\n",
        "import cv2\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "image = np.array(image)\n",
        "low_threshold = 100\n",
        "high_threshold = 200\n",
        "image = cv2.Canny(image, low_threshold, high_threshold)\n",
        "image = image[:, :, None]\n",
        "image = np.concatenate([image, image, image], axis=2)\n",
        "canny_image = Image.fromarray(image)\n",
        "image = image_grid([image_raw, canny_image], 1, 2)\n",
        "image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title ControlNet #Canny - Step 3\n",
        "gc.collect()\n",
        "\n",
        "prompt = \"cute duck\"\n",
        "negative_prompt = \"blurry, ugly\"\n",
        "\n",
        "num_samples = jax.device_count()\n",
        "\n",
        "prompt_n = num_samples * [prompt]\n",
        "prompt_ids = pipeline.prepare_inputs(prompt_n)\n",
        "prompt_ids = shard(prompt_ids)\n",
        "\n",
        "negative_prompt_n = num_samples * [negative_prompt]\n",
        "negative_prompt_ids = pipeline.prepare_inputs(negative_prompt_n)\n",
        "negative_prompt_ids = shard(negative_prompt_ids)\n",
        "\n",
        "real_seed = random.randint(0, 2147483647)\n",
        "prng_seed = jax.random.PRNGKey(real_seed)\n",
        "prng_seed = jax.random.split(prng_seed, jax.device_count())\n",
        "\n",
        "images = pipeline(prompt_ids, canny_image, params, prng_seed, neg_prompt_ids=negative_prompt_ids, num_inference_steps=50, height=512, width=512, guidance_scale=7, jit=True).images\n",
        "images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))\n",
        "\n",
        "image = image_grid(images, 2, 4)\n",
        "image"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}