pytorch_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/pytorch_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 -q torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 torchtext==0.14.1 torchdata==0.5.1 --extra-index-url https://download.pytorch.org/whl/cu116 -U\n",
        "!pip install -U git+https://github.com/huggingface/diffusers\n",
        "!pip install -U transformers piexif fold_to_ascii ftfy\n",
        "!pip install -q https://github.com/camenduru/stable-diffusion-webui-colab/releases/download/0.0.15/xformers-0.0.15.dev0+189828c.d20221207-cp38-cp38-linux_x86_64.whl\n",
        "\n",
        "import torch, os, gc, requests, json, piexif\n",
        "from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler\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",
        "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 = \"stabilityai/stable-diffusion-2\" #@param [\"stabilityai/stable-diffusion-2\", \"stabilityai/stable-diffusion-2-base\", \"stabilityai/stable-diffusion-2-inpainting\", \"stabilityai/stable-diffusion-2-depth\"] {allow-input: true}\n",
        "\n",
        "is_tile = False #@param {type: 'boolean'}\n",
        "if(is_tile):\n",
        "  def patch_conv(cls):\n",
        "    init = cls.__init__\n",
        "    def __init__(self, *args, **kwargs):\n",
        "      return init(self, *args, **kwargs, padding_mode='circular')\n",
        "    cls.__init__ = __init__\n",
        "  patch_conv(torch.nn.Conv2d)\n",
        "\n",
        "scheduler = EulerDiscreteScheduler.from_pretrained(model_folder, subfolder=\"scheduler\")\n",
        "pipe = StableDiffusionPipeline.from_pretrained(model_folder, scheduler=scheduler).to(\"cuda\")\n",
        "pipe.safety_checker = lambda images, clip_input: (images, False)\n",
        "pipe.enable_xformers_memory_efficient_attention()\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",
        "    with torch.autocast(\"cuda\"):\n",
        "      images = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, height=height, width=width, guidance_scale=guidance_scale).images\n",
        "    image = images[0]\n",
        "    real_seed = torch.cuda.initial_seed()\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 = \"\" #@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  = 768 #@param {type: 'integer'}\n",
        "height  = 768 #@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": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}