flax_stable_diffusion_2.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_2.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ymcDoPQn41U4"
      },
      "outputs": [],
      "source": [
        "save_to = \"here\" #@param [\"here\", \"discord\", \"gdrive\"]\n",
        "if save_to == \"gdrive\":\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/gdrive')\n",
        "\n",
        "!pip install -U jax==0.3.25 jaxlib==0.3.25 flax==0.6.2 transformers piexif fold_to_ascii ftfy\n",
        "!pip install -U diffusers==0.10.0\n",
        "\n",
        "import jax.tools.colab_tpu\n",
        "jax.tools.colab_tpu.setup_tpu('tpu_driver_nightly')\n",
        "\n",
        "import jax, random, os, gc, requests, json, piexif\n",
        "import numpy as np\n",
        "from flax.jax_utils import replicate\n",
        "from flax.training.common_utils import shard\n",
        "from diffusers import FlaxStableDiffusionPipeline\n",
        "from PIL import Image\n",
        "from PIL.PngImagePlugin import PngInfo\n",
        "from fold_to_ascii import fold\n",
        "metadata = PngInfo()\n",
        "\n",
        "def closestNumber(n, m):\n",
        "    q = int(n / m)\n",
        "    n1 = m * q\n",
        "    if (n * m) > 0:\n",
        "        n2 = m * (q + 1)\n",
        "    else:\n",
        "        n2 = m * (q - 1)\n",
        "    if abs(n - n1) < abs(n - n2):\n",
        "        return n1\n",
        "    return n2\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\n",
        "\n",
        "folder_max_files = 500 #@param {type: 'integer'}\n",
        "root_folder = 'ai_images' #@param {type: 'string'}\n",
        "if save_to == \"gdrive\":\n",
        "  root_folder = f\"/content/gdrive/MyDrive/{root_folder}\"\n",
        "\n",
        "if os.path.exists(f\"{root_folder}\") == False:\n",
        "    os.mkdir(f\"{root_folder}\")\n",
        "image_folder = max([int(f) for f in os.listdir(f\"{root_folder}\")], default=0)\n",
        "if os.path.exists(f\"{root_folder}/{image_folder:04}\") == False:\n",
        "    os.mkdir(f\"{root_folder}/{image_folder:04}\")\n",
        "name = max([int(f[: f.index(\".\")]) for f in os.listdir(f\"{root_folder}/{image_folder:04}\")],default=0,)\n",
        "\n",
        "model_folder = \"flax/stable-diffusion-2\" #@param [\"flax/stable-diffusion-2\", \"flax/stable-diffusion-2-base\"] {allow-input: true}\n",
        "revision = \"main\"\n",
        "pipe, params = FlaxStableDiffusionPipeline.from_pretrained(model_folder, revision=revision, from_pt=False, dtype=jax.numpy.bfloat16, safety_checker=None)\n",
        "params = replicate(params)\n",
        "\n",
        "def generate(discord_token, discord_channel_id, discord_user, by, num_inference_steps, guidance_scale, sampler, width, height, prompt, negative_prompt, suffix, image_folder, name):\n",
        "    width = closestNumber(width, 8)\n",
        "    height = closestNumber(height, 8)\n",
        "    metadata.add_text(\"Prompt\", f\"{prompt}\")\n",
        "    metadata.add_text(\"by\", f\"{by}\")\n",
        "    gc.collect()\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",
        "    num_samples = jax.device_count()\n",
        "    prompt_n = num_samples * [prompt]\n",
        "    prompt_ids = pipe.prepare_inputs(prompt_n)\n",
        "    prompt_ids = shard(prompt_ids)\n",
        "    negative_prompt_n = num_samples * [negative_prompt]\n",
        "    negative_prompt_ids = pipe.prepare_inputs(negative_prompt_n)\n",
        "    negative_prompt_ids = shard(negative_prompt_ids)\n",
        "    images = pipe(prompt_ids, params, prng_seed, neg_prompt_ids=negative_prompt_ids, num_inference_steps=num_inference_steps, height=height, width=width, guidance_scale=guidance_scale, jit=True).images\n",
        "    images = pipe.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))\n",
        "    image = image_grid(images, 2, 4)\n",
        "    if(suffix == 'png'):\n",
        "      image.save(f\"{root_folder}/{image_folder:04}/{name:04}.{suffix}\", pnginfo=metadata)\n",
        "    else:\n",
        "      zeroth_ifd = {piexif.ImageIFD.ImageDescription: f\"{fold(prompt)}\", piexif.ImageIFD.Make: f\"{fold(by)}\", piexif.ImageIFD.Model: f\"{model_folder}\"}\n",
        "      exif_dict = {\"0th\": zeroth_ifd}\n",
        "      exif_bytes = piexif.dump(exif_dict)\n",
        "      image.save(f\"{root_folder}/{image_folder:04}/{name:04}.{suffix}\", \"JPEG\", quality=70, exif=exif_bytes)\n",
        "    files = {f\"{image_folder:04}_{name:04}.{suffix}\": open(f\"{root_folder}/{image_folder:04}/{name:04}.{suffix}\", \"rb\").read()}\n",
        "    if save_to == \"discord\":\n",
        "      payload = {\"content\": f\"{prompt}\\nNegative prompt: {negative_prompt}\\nSteps: {num_inference_steps}, Sampler: {sampler}, CFG scale: {guidance_scale}, Seed: {real_seed}, Size: {width}x{height}, Model folder: {model_folder} - {discord_user}\"}\n",
        "      requests.post(f\"https://discord.com/api/v9/channels/{discord_channel_id}/messages\", data=payload, headers={\"authorization\": f\"Bot {discord_token}\"}, files=files)\n",
        "      os.remove(f\"{root_folder}/{image_folder:04}/{name:04}.{suffix}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "DDUev8TOMxJ4"
      },
      "outputs": [],
      "source": [
        "discord_token = \"token\" #@param {type: 'string'}\n",
        "discord_channel_id = 0 #@param {type: 'integer'}\n",
        "prompt = \"duck\" #@param {type: 'string'}\n",
        "negative_prompt = \"\" #@param {type: 'string'}\n",
        "width  = 512 #@param {type: 'integer'}\n",
        "height  = 512 #@param {type: 'integer'}\n",
        "guidance_scale = 7.5 #@param {type: 'number'}\n",
        "num_inference_steps = 50 #@param {type: 'integer'}\n",
        "suffix = \"jpg\" #@param [\"jpg\", \"png\"]\n",
        "by = \"camenduru\" #@param {type: 'string'}\n",
        "template = {\n",
        "    \"discord_token\": discord_token,\n",
        "    \"discord_channel_id\": discord_channel_id,\n",
        "    \"by\": by,\n",
        "    \"num_inference_steps\": num_inference_steps,\n",
        "    \"guidance_scale\": guidance_scale,\n",
        "    \"sampler\": \"PLMS\",\n",
        "    \"width\": width,\n",
        "    \"height\": height,\n",
        "    \"prompt\": prompt,\n",
        "    \"negative_prompt\": negative_prompt,\n",
        "    \"suffix\": suffix\n",
        "  }\n",
        "with open(\"template.json\", \"w\") as outfile:\n",
        "    outfile.write(json.dumps(template))\n",
        "\n",
        "is_loop = False #@param {type:\"boolean\"}\n",
        "\n",
        "if is_loop:\n",
        "  while True:\n",
        "      if name < folder_max_files:\n",
        "          with open(\"template.json\", \"r\") as file:\n",
        "              prompts = file.readlines()\n",
        "          for prompt in prompts:\n",
        "              d = json.loads(prompt)\n",
        "              name += 1\n",
        "              generate(d[\"discord_token\"], d[\"discord_channel_id\"], \"camenduru\", d[\"by\"], d[\"num_inference_steps\"], d[\"guidance_scale\"], d[\"sampler\"], d[\"width\"], d[\"height\"], d[\"prompt\"], d[\"negative_prompt\"], d[\"suffix\"], image_folder, name)\n",
        "      else:\n",
        "          image_folder += 1\n",
        "          if os.path.exists(f\"{root_folder}/{image_folder:04}\") == False:\n",
        "              os.mkdir(f\"{root_folder}/{image_folder:04}\")\n",
        "          name = 0\n",
        "else:\n",
        "  if name < folder_max_files:\n",
        "      with open(\"template.json\", \"r\") as file:\n",
        "          prompts = file.readlines()\n",
        "      for prompt in prompts:\n",
        "          d = json.loads(prompt)\n",
        "          name += 1\n",
        "          generate(d[\"discord_token\"], d[\"discord_channel_id\"], \"camenduru\", d[\"by\"], d[\"num_inference_steps\"], d[\"guidance_scale\"], d[\"sampler\"], d[\"width\"], d[\"height\"], d[\"prompt\"], d[\"negative_prompt\"], d[\"suffix\"], image_folder, name)\n",
        "  else:\n",
        "      image_folder += 1\n",
        "      if os.path.exists(f\"{root_folder}/{image_folder:04}\") == False:\n",
        "          os.mkdir(f\"{root_folder}/{image_folder:04}\")\n",
        "      name = 0"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}