fast_stable_diffusion_AUTOMATIC1111.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# **Colab From https://github.com/TheLastBen/fast-stable-diffusion, if you face any issues, feel free to discuss them.**\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "47kV9o1Ni8GH"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y9EBc437WDOs"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # Installing AUTOMATIC1111 repo\n",
        "%%capture\n",
        "%cd /content/gdrive/MyDrive/\n",
        "%mkdir  sd\n",
        "%cd sd\n",
        "!git clone https://github.com/CompVis/stable-diffusion\n",
        "!git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui\n",
        "%cd /content/gdrive/MyDrive/sd/stable-diffusion-webui/\n",
        "!mkdir -p cache/{huggingface,torch}\n",
        "%cd /content/\n",
        "!ln -s /content/gdrive/MyDrive/sd/stable-diffusion-webui/cache/huggingface ../root/.cache/\n",
        "!ln -s /content/gdrive/MyDrive/sd/stable-diffusion-webui/cache/torch ../root/.cache/\n"       
      ],
      "metadata": {
        "id": "CFWtw-6EPrKi",
        "cellView": "form"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # Set paths\n",
        "%%writefile /content/gdrive/MyDrive/sd/stable-diffusion-webui/modules/paths.py\n",
        "import argparse\n",
        "import os\n",
        "import sys\n",
        "\n",
        "script_path = '/content/gdrive/MyDrive/sd/stable-diffusion-webui'\n",
        "models_path = os.path.join(script_path, \"models\")\n",
        "sys.path.insert(0, script_path)\n",
        "\n",
        "# search for directory of stable diffusion in following places\n",
        "sd_path = None\n",
        "possible_sd_paths = [os.path.join(script_path, '/content/gdrive/MyDrive/sd/stable-diffusion'), '.', os.path.dirname(script_path)]\n",
        "for possible_sd_path in possible_sd_paths:\n",
        "    if os.path.exists(os.path.join(possible_sd_path, 'ldm/models/diffusion/ddpm.py')):\n",
        "        sd_path = os.path.abspath(possible_sd_path)\n",
        "\n",
        "assert sd_path is not None, \"Couldn't find Stable Diffusion in any of: \" + str(possible_sd_paths)\n",
        "\n",
        "path_dirs = [\n",
        "    (sd_path, 'ldm', 'Stable Diffusion', []),\n",
        "    (os.path.join(sd_path, 'src/taming-transformers'), 'taming', 'Taming Transformers', []),\n",
        "    (os.path.join(sd_path, 'src/codeformer'), 'inference_codeformer.py', 'CodeFormer', []),\n",
        "    (os.path.join(sd_path, 'src/blip'), 'models/blip.py', 'BLIP', []),\n",
        "    (os.path.join(sd_path, 'src/latent-diffusion'), 'LDSR.py', 'LDSR', []),\n",
        "    (os.path.join(sd_path, 'src/k-diffusion'), 'k_diffusion/sampling.py', 'k_diffusion', [\"atstart\"]),\n",
        "]\n",
        "\n",
        "paths = {}\n",
        "\n",
        "for d, must_exist, what, options in path_dirs:\n",
        "    must_exist_path = os.path.abspath(os.path.join(script_path, d, must_exist))\n",
        "    if not os.path.exists(must_exist_path):\n",
        "        print(f\"Warning: {what} not found at path {must_exist_path}\", file=sys.stderr)\n",
        "    else:\n",
        "        d = os.path.abspath(d)\n",
        "        if \"atstart\" in options:\n",
        "            sys.path.insert(0, d)\n",
        "        else:\n",
        "            sys.path.append(d)\n",
        "        paths[what] = d"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Vl9UjD6A6rmp",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f6d6e084-0700-44db-9df7-46f19829ec7e"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Overwriting paths.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from IPython.display import clear_output\n",
        "\n",
        "#@markdown # Model Download/Load\n",
        "token = \"\" #@param {type:\"string\"}\n",
        "#@markdown \n",
        "Path_to_trained_model = \"\" #@param {type:\"string\"}\n",
        "#@markdown Insert the full path of your trained model (eg: /content/gdrive/MyDrive/zarathustra.ckpt) and it will automatically be placed in the right place, otherwise, leave it EMPTY (make sure there are no spaces in the path)\n",
        "if (Path_to_trained_model !=''):\n",
        "  if os.path.exists(str(Path_to_trained_model)):\n",
        "    clear_output()\n",
        "    !mv -f $Path_to_trained_model '/content/gdrive/MyDrive/sd/stable-diffusion-webui/models/Stable-diffusion/model.ckpt'\n",
        "    if os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion-webui/models/Stable-diffusion/model.ckpt'):\n",
        "      print('Model placed in the right directory')\n",
        "    else:\n",
        "      print('Something went wrong')\n",
        "  else:\n",
        "    print('Wrong path, use the colab file explorer to copy the path')\n",
        "\n",
        "else:\n",
        "\n",
        "  if token == \"\" and not os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion-webui/models/Stable-diffusion/model.ckpt'):\n",
        "    token=input(\"Insert your huggingface token :\")\n",
        "    %cd /content/\n",
        "    !git init\n",
        "    !git lfs install --system --skip-repo\n",
        "    !git remote add -f origin \"https://USER:{token}@huggingface.co/CompVis/stable-diffusion-v-1-4-original\"\n",
        "    !git config core.sparsecheckout true\n",
        "    !echo \"sd-v1-4.ckpt\" > .git/info/sparse-checkout\n",
        "    !git pull origin main\n",
        "    !mv '/content/sd-v1-4.ckpt' '/content/gdrive/MyDrive/sd/stable-diffusion-webui/models/Stable-diffusion/model.ckpt'\n",
        "    if os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion-webui/models/Stable-diffusion/model.ckpt'):\n",
        "      clear_output()\n",
        "      print(\"Model successfully downloaded\")  \n",
        "\n",
        "  elif not os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion-webui/models/Stable-diffusion/model.ckpt'):\n",
        "        %cd /content/\n",
        "        !git init\n",
        "        !git lfs install --system --skip-repo\n",
        "        !git remote add -f origin \"https://USER:{token}@huggingface.co/CompVis/stable-diffusion-v-1-4-original\"\n",
        "        !git config core.sparsecheckout true\n",
        "        !echo \"sd-v1-4.ckpt\" > .git/info/sparse-checkout\n",
        "        !git pull origin main\n",
        "        !mv '/content/sd-v1-4.ckpt' '/content/gdrive/MyDrive/sd/stable-diffusion-webui/models/Stable-diffusion/model.ckpt'\n",
        "        if os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion-webui/models/Stable-diffusion/model.ckpt'):\n",
        "          clear_output()\n",
        "          print(\"Model successfully downloaded\")  \n",
        "\n",
        "  else:\n",
        "      print(\"Model already exists\")\n",
        "\n",
        "  if os.path.exists('/content/.git'):\n",
        "    !rm -r /content/.git"
      ],
      "metadata": {
        "id": "p4wj_txjP3TC",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "ZGV_5H4xrOSp",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@markdown # Installing Requirements\n",
        "%%capture\n",
        "import os\n",
        "if not os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion/src/k-diffusion/k_diffusion'):\n",
        "  %cd /content/gdrive/MyDrive/sd/stable-diffusion/\n",
        "  !pip install -e git+https://github.com/CompVis/taming-transformers#egg=taming-transformers\n",
        "  !pip install -e git+https://github.com/openai/CLIP#egg=clip\n",
        "  !pip install -e git+https://github.com/TencentARC/GFPGAN#egg=gfpgan\n",
        "  !pip install -e git+https://github.com/salesforce/BLIP#egg=blip\n",
        "  !pip install -e git+https://github.com/sczhou/CodeFormer#egg=codeformer\n",
        "  !pip install -e git+https://github.com/xinntao/Real-ESRGAN#egg=realesrgan\n",
        "  !pip install -e git+https://github.com/crowsonkb/k-diffusion.git#egg=k_diffusion\n",
        "  !pip install -e git+https://github.com/Hafiidz/latent-diffusion#egg=latent-diffusion\n",
        "!pip install git+https://github.com/openai/CLIP\n",
        "!pip install realesrgan\n",
        "!pip install torchdiffeq \n",
        "!pip install resize_right\n",
        "!pip install einops\n",
        "!pip install jsonmerge\n",
        "!pip install clean-fid\n",
        "!pip install gfpgan\n",
        "!pip install addict\n",
        "!pip install future\n",
        "!pip install lmdb\n",
        "!pip install pyyaml\n",
        "!pip install requests\n",
        "!pip install scipy\n",
        "!pip install tb-nightly\n",
        "!pip install tqdm\n",
        "!pip install yapf\n",
        "!pip install lpips\n",
        "!pip install gdown\n",
        "!pip install timm==0.4.12\n",
        "!pip install fairscale==0.4.4\n",
        "!pip install pycocoevalcap\n",
        "!pip install scikit-image\n",
        "!pip install numpy==1.21.6\n",
        "!pip install albumentations==0.4.3\n",
        "!pip install diffusers==0.3.0\n",
        "!pip install facexlib>=0.2.3\n",
        "!pip install gradio==3.4b3\n",
        "!pip install imageio-ffmpeg==0.4.2\n",
        "!pip install imageio==2.9.0\n",
        "!pip install kornia==0.6\n",
        "!pip install omegaconf==2.1.1\n",
        "!pip install opencv-python-headless==4.6.0.66\n",
        "!pip install piexif==1.1.3\n",
        "!pip install pudb==2019.2\n",
        "!pip install pynvml==11.4.1\n",
        "!pip install fonts font-roboto\n",
        "!pip install python-slugify>=6.1.2\n",
        "!pip install pytorch-lightning\n",
        "!pip install torch-fidelity==0.3.0\n",
        "!pip install transformers==4.19.2\n",
        "!pip install triton==2.0.0.dev20220701"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # txt2mask by ThereforeGames\n",
        "%%capture\n",
        "import os\n",
        "\n",
        "if not os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion-webui/scripts/txt2mask.py'):\n",
        "  %cd /content/\n",
        "  !git clone https://github.com/ThereforeGames/txt2mask\n",
        "  !mv /content/txt2mask/scripts/* /content/gdrive/MyDrive/sd/stable-diffusion-webui/scripts\n",
        "  %cd /content/txt2mask\n",
        "  !rmdir 'scripts'\n",
        "  !mv /content/txt2mask/* /content/gdrive/MyDrive/sd/stable-diffusion-webui/"
      ],
      "metadata": {
        "id": "WhlZBT7RDSHK",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # Installing xformers\n",
        "\n",
        "from IPython.display import clear_output\n",
        "import time\n",
        "from IPython.display import HTML\n",
        "from subprocess import getoutput\n",
        "import os\n",
        "\n",
        "if not os.path.exists('/content/gdrive/MyDrive/sd/stable-diffusion-webui/k_diffusion'):\n",
        "  !cp -R '/content/gdrive/MyDrive/sd/stable-diffusion/src/k-diffusion/k_diffusion' '/content/gdrive/MyDrive/sd/stable-diffusion-webui/'\n",
        "  !cp -R '/content/gdrive/MyDrive/sd/stable-diffusion/ldm' '/content/gdrive/MyDrive/sd/stable-diffusion-webui/'\n",
        "\n",
        "s = getoutput('nvidia-smi')\n",
        "if 'T4' in s:\n",
        "  gpu = 'T4'\n",
        "elif 'P100' in s:\n",
        "  gpu = 'P100'\n",
        "elif 'V100' in s:\n",
        "  gpu = 'V100'\n",
        "elif 'A100' in s:\n",
        "  gpu = 'A100'\n",
        "\n",
        "while True:\n",
        "    try: \n",
        "        gpu=='T4'or gpu=='P100'or gpu=='V100'or gpu=='A100'\n",
        "        break\n",
        "    except:\n",
        "        pass\n",
        "    print('it seems that your GPU is not supported at the moment')\n",
        "    time.sleep(5)\n",
        "\n",
        "if (gpu=='T4'):\n",
        "  %pip install https://github.com/TheLastBen/fast-stable-diffusion/raw/main/precompiled/T4/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "  \n",
        "elif (gpu=='P100'):\n",
        "  %pip install https://github.com/TheLastBen/fast-stable-diffusion/raw/main/precompiled/P100/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "\n",
        "elif (gpu=='V100'):\n",
        "  %pip install https://github.com/TheLastBen/fast-stable-diffusion/raw/main/precompiled/V100/xformers-0.0.13.dev0-py3-none-any.whl\n",
        "\n",
        "elif (gpu=='A100'):\n",
        "  %pip install https://github.com/TheLastBen/fast-stable-diffusion/raw/main/precompiled/A100/xformers-0.0.13.dev0-py3-none-any.whl  \n",
        "\n",
        "clear_output()\n",
        "print('DONE !')\n"
      ],
      "metadata": {
        "id": "a---cT2rwUQj",
        "cellView": "form"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # Patching attention.py\n",
        "%%writefile /content/gdrive/MyDrive/sd/stable-diffusion-webui/ldm/modules/attention.py\n",
        "import gc\n",
        "from inspect import isfunction\n",
        "import math\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch import nn, einsum\n",
        "from einops import rearrange, repeat\n",
        "import os\n",
        "from typing import Any, Optional\n",
        "import xformers\n",
        "import xformers.ops\n",
        "\n",
        "\n",
        "from ldm.modules.diffusionmodules.util import checkpoint\n",
        "\n",
        "\n",
        "def exists(val):\n",
        "    return val is not None\n",
        "\n",
        "\n",
        "def uniq(arr):\n",
        "    return{el: True for el in arr}.keys()\n",
        "\n",
        "\n",
        "def default(val, d):\n",
        "    if exists(val):\n",
        "        return val\n",
        "    return d() if isfunction(d) else d\n",
        "\n",
        "\n",
        "def max_neg_value(t):\n",
        "    return -torch.finfo(t.dtype).max\n",
        "\n",
        "\n",
        "def init_(tensor):\n",
        "    dim = tensor.shape[-1]\n",
        "    std = 1 / math.sqrt(dim)\n",
        "    tensor.uniform_(-std, std)\n",
        "    return tensor\n",
        "\n",
        "\n",
        "\n",
        "# feedforward\n",
        "class GEGLU(nn.Module):\n",
        "    def __init__(self, dim_in, dim_out):\n",
        "        super().__init__()\n",
        "        self.proj = nn.Linear(dim_in, dim_out * 2)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x, gate = self.proj(x).chunk(2, dim=-1)\n",
        "        return x * F.gelu(gate)\n",
        "\n",
        "\n",
        "class FeedForward(nn.Module):\n",
        "    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):\n",
        "        super().__init__()\n",
        "        inner_dim = int(dim * mult)\n",
        "        dim_out = default(dim_out, dim)\n",
        "        project_in = nn.Sequential(\n",
        "            nn.Linear(dim, inner_dim),\n",
        "            nn.GELU()\n",
        "        ) if not glu else GEGLU(dim, inner_dim)\n",
        "\n",
        "        self.net = nn.Sequential(\n",
        "            project_in,\n",
        "            nn.Dropout(dropout),\n",
        "            nn.Linear(inner_dim, dim_out)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.net(x)\n",
        "\n",
        "\n",
        "def zero_module(module):\n",
        "    \"\"\"\n",
        "    Zero out the parameters of a module and return it.\n",
        "    \"\"\"\n",
        "    for p in module.parameters():\n",
        "        p.detach().zero_()\n",
        "    return module\n",
        "\n",
        "\n",
        "def Normalize(in_channels):\n",
        "    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)\n",
        "\n",
        "\n",
        "class LinearAttention(nn.Module):\n",
        "    def __init__(self, dim, heads=4, dim_head=32):\n",
        "        super().__init__()\n",
        "        self.heads = heads\n",
        "        hidden_dim = dim_head * heads\n",
        "        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)\n",
        "        self.to_out = nn.Conv2d(hidden_dim, dim, 1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        b, c, h, w = x.shape\n",
        "        qkv = self.to_qkv(x)\n",
        "        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)\n",
        "        k = k.softmax(dim=-1)\n",
        "        context = torch.einsum('bhdn,bhen->bhde', k, v)\n",
        "        out = torch.einsum('bhde,bhdn->bhen', context, q)\n",
        "        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)\n",
        "        return self.to_out(out)\n",
        "\n",
        "\n",
        "class SpatialSelfAttention(nn.Module):\n",
        "    def __init__(self, in_channels):\n",
        "        super().__init__()\n",
        "        self.in_channels = in_channels\n",
        "\n",
        "        self.norm = Normalize(in_channels)\n",
        "        self.q = torch.nn.Conv2d(in_channels,\n",
        "                                 in_channels,\n",
        "                                 kernel_size=1,\n",
        "                                 stride=1,\n",
        "                                 padding=0)\n",
        "        self.k = torch.nn.Conv2d(in_channels,\n",
        "                                 in_channels,\n",
        "                                 kernel_size=1,\n",
        "                                 stride=1,\n",
        "                                 padding=0)\n",
        "        self.v = torch.nn.Conv2d(in_channels,\n",
        "                                 in_channels,\n",
        "                                 kernel_size=1,\n",
        "                                 stride=1,\n",
        "                                 padding=0)\n",
        "        self.proj_out = torch.nn.Conv2d(in_channels,\n",
        "                                        in_channels,\n",
        "                                        kernel_size=1,\n",
        "                                        stride=1,\n",
        "                                        padding=0)\n",
        "\n",
        "    def forward(self, x):\n",
        "        h_ = x\n",
        "        h_ = self.norm(h_)\n",
        "        q = self.q(h_)\n",
        "        k = self.k(h_)\n",
        "        v = self.v(h_)\n",
        "\n",
        "        # compute attention\n",
        "        b,c,h,w = q.shape\n",
        "        q = rearrange(q, 'b c h w -> b (h w) c')\n",
        "        k = rearrange(k, 'b c h w -> b c (h w)')\n",
        "        w_ = torch.einsum('bij,bjk->bik', q, k)\n",
        "\n",
        "        w_ = w_ * (int(c)**(-0.5))\n",
        "        w_ = torch.nn.functional.softmax(w_, dim=2)\n",
        "\n",
        "        # attend to values\n",
        "        v = rearrange(v, 'b c h w -> b c (h w)')\n",
        "        w_ = rearrange(w_, 'b i j -> b j i')\n",
        "        h_ = torch.einsum('bij,bjk->bik', v, w_)\n",
        "        h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)\n",
        "        h_ = self.proj_out(h_)\n",
        "\n",
        "        return x+h_\n",
        "\n",
        "\n",
        "class CrossAttention(nn.Module):\n",
        "    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):\n",
        "        super().__init__()\n",
        "        inner_dim = dim_head * heads\n",
        "        context_dim = default(context_dim, query_dim)\n",
        "\n",
        "        self.scale = dim_head ** -0.5\n",
        "        self.heads = heads\n",
        "\n",
        "        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)\n",
        "        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)\n",
        "        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)\n",
        "        self.to_out = nn.Sequential(\n",
        "            nn.Linear(inner_dim, query_dim),\n",
        "            nn.Dropout(dropout)\n",
        "        )\n",
        "\n",
        "    def forward(self, x, context=None, mask=None):\n",
        "        h = self.heads\n",
        "\n",
        "        q_in = self.to_q(x)\n",
        "        context = default(context, x)\n",
        "        k_in = self.to_k(context)\n",
        "        v_in = self.to_v(context)\n",
        "        del context, x\n",
        "\n",
        "        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))\n",
        "        del q_in, k_in, v_in\n",
        "\n",
        "        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)\n",
        "\n",
        "        stats = torch.cuda.memory_stats(q.device)\n",
        "        mem_active = stats['active_bytes.all.current']\n",
        "        mem_reserved = stats['reserved_bytes.all.current']\n",
        "        mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())\n",
        "        mem_free_torch = mem_reserved - mem_active\n",
        "        mem_free_total = mem_free_cuda + mem_free_torch\n",
        "\n",
        "        gb = 1024 ** 3\n",
        "        tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()\n",
        "        modifier = 3 if q.element_size() == 2 else 2.5\n",
        "        mem_required = tensor_size * modifier\n",
        "        steps = 1\n",
        "\n",
        "\n",
        "        if mem_required > mem_free_total:\n",
        "            steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))\n",
        "            # print(f\"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB \"\n",
        "            #      f\"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}\")\n",
        "\n",
        "        if steps > 64:\n",
        "            max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64\n",
        "            raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '\n",
        "                               f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free')\n",
        "\n",
        "        slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]\n",
        "        for i in range(0, q.shape[1], slice_size):\n",
        "            end = i + slice_size\n",
        "            s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale\n",
        "\n",
        "            s2 = s1.softmax(dim=-1, dtype=q.dtype)\n",
        "            del s1\n",
        "\n",
        "            r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)\n",
        "            del s2\n",
        "\n",
        "        del q, k, v\n",
        "\n",
        "        r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)\n",
        "        del r1\n",
        "\n",
        "        return self.to_out(r2)\n",
        "\n",
        "\n",
        "class BasicTransformerBlock(nn.Module):\n",
        "    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True):\n",
        "        super().__init__()\n",
        "        AttentionBuilder = MemoryEfficientCrossAttention        \n",
        "        self.attn1 = AttentionBuilder(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout)  # is a self-attention\n",
        "        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)\n",
        "        self.attn2 = AttentionBuilder(query_dim=dim, context_dim=context_dim,\n",
        "                                    heads=n_heads, dim_head=d_head, dropout=dropout)  # is self-attn if context is none\n",
        "        self.norm1 = nn.LayerNorm(dim)\n",
        "        self.norm2 = nn.LayerNorm(dim)\n",
        "        self.norm3 = nn.LayerNorm(dim)\n",
        "        self.checkpoint = checkpoint\n",
        "        \n",
        "    def _set_attention_slice(self, slice_size):\n",
        "        self.attn1._slice_size = slice_size\n",
        "        self.attn2._slice_size = slice_size\n",
        "\n",
        "    def forward(self, hidden_states, context=None):\n",
        "        hidden_states = hidden_states.contiguous() if hidden_states.device.type == \"mps\" else hidden_states\n",
        "        hidden_states = self.attn1(self.norm1(hidden_states)) + hidden_states\n",
        "        hidden_states = self.attn2(self.norm2(hidden_states), context=context) + hidden_states\n",
        "        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states\n",
        "        return hidden_states        \n",
        "\n",
        "    # def forward(self, x, context=None):\n",
        "        # return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)\n",
        "\n",
        "    # def _forward(self, x, context=None):\n",
        "        # x = self.attn1(self.norm1(x)) + x\n",
        "        # x = self.attn2(self.norm2(x), context=context) + x\n",
        "        # x = self.ff(self.norm3(x)) + x\n",
        "        # return x\n",
        "\n",
        "class MemoryEfficientCrossAttention(nn.Module):\n",
        "    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):\n",
        "        super().__init__()\n",
        "        inner_dim = dim_head * heads\n",
        "        context_dim = default(context_dim, query_dim)\n",
        "\n",
        "        self.scale = dim_head**-0.5\n",
        "        self.heads = heads\n",
        "        self.dim_head = dim_head\n",
        "\n",
        "        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)\n",
        "        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)\n",
        "        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)\n",
        "\n",
        "        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))\n",
        "        self.attention_op: Optional[Any] = None\n",
        "\n",
        "    def _maybe_init(self, x):\n",
        "        \"\"\"\n",
        "        Initialize the attention operator, if required We expect the head dimension to be exposed here, meaning that x\n",
        "        : B, Head, Length\n",
        "        \"\"\"\n",
        "        if self.attention_op is not None:\n",
        "            return\n",
        "\n",
        "        _, M, K = x.shape\n",
        "        try:\n",
        "            self.attention_op = xformers.ops.AttentionOpDispatch(\n",
        "                dtype=x.dtype,\n",
        "                device=x.device,\n",
        "                k=K,\n",
        "                attn_bias_type=type(None),\n",
        "                has_dropout=False,\n",
        "                kv_len=M,\n",
        "                q_len=M,\n",
        "            ).op\n",
        "\n",
        "        except NotImplementedError as err:\n",
        "            raise NotImplementedError(f\"Please install xformers with the flash attention / cutlass components.\\n{err}\")\n",
        "\n",
        "    def forward(self, x, context=None, mask=None):\n",
        "\n",
        "\n",
        "        q = self.to_q(x)\n",
        "        context = default(context, x)\n",
        "        k = self.to_k(context)\n",
        "        v = self.to_v(context)\n",
        "        \n",
        "\n",
        "\n",
        "        b, _, _ = q.shape\n",
        "        q, k, v = map(\n",
        "            lambda t: t.unsqueeze(3)\n",
        "            .reshape(b, t.shape[1], self.heads, self.dim_head)\n",
        "            .permute(0, 2, 1, 3)\n",
        "            .reshape(b * self.heads, t.shape[1], self.dim_head)\n",
        "            .contiguous(),\n",
        "            (q, k, v),\n",
        "        )\n",
        "\n",
        "        # init the attention op, if required, using the proper dimensions\n",
        "        self._maybe_init(q)\n",
        "\n",
        "        # actually compute the attention, what we cannot get enough of\n",
        "        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)\n",
        "\n",
        "        # TODO: Use this directly in the attention operation, as a bias\n",
        "        if exists(mask):\n",
        "            raise NotImplementedError\n",
        "        out = (\n",
        "            out.unsqueeze(0)\n",
        "            .reshape(b, self.heads, out.shape[1], self.dim_head)\n",
        "            .permute(0, 2, 1, 3)\n",
        "            .reshape(b, out.shape[1], self.heads * self.dim_head)\n",
        "        )\n",
        "\n",
        "        stats = torch.cuda.memory_stats(q.device)\n",
        "        mem_active = stats['active_bytes.all.current']\n",
        "        mem_reserved = stats['reserved_bytes.all.current']\n",
        "        mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())\n",
        "        mem_free_torch = mem_reserved - mem_active\n",
        "        mem_free_total = mem_free_cuda + mem_free_torch\n",
        "\n",
        "        gb = 1024 ** 3\n",
        "        tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()\n",
        "        modifier = 3 if q.element_size() == 2 else 2.5\n",
        "        mem_required = tensor_size * modifier\n",
        "        steps = 1\n",
        "\n",
        "\n",
        "        if mem_required > mem_free_total:\n",
        "            steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))\n",
        "            # print(f\"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB \"\n",
        "            #      f\"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}\")\n",
        "\n",
        "        if steps > 64:\n",
        "            max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64\n",
        "            raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '\n",
        "                               f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free')\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "        return self.to_out(out)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class SpatialTransformer(nn.Module):\n",
        "    \"\"\"\n",
        "    Transformer block for image-like data.\n",
        "    First, project the input (aka embedding)\n",
        "    and reshape to b, t, d.\n",
        "    Then apply standard transformer action.\n",
        "    Finally, reshape to image\n",
        "    \"\"\"\n",
        "    def __init__(self, in_channels, n_heads, d_head,\n",
        "                 depth=1, dropout=0., context_dim=None):\n",
        "        super().__init__()\n",
        "        self.in_channels = in_channels\n",
        "        inner_dim = n_heads * d_head\n",
        "        self.norm = Normalize(in_channels)\n",
        "\n",
        "        self.proj_in = nn.Conv2d(in_channels,\n",
        "                                 inner_dim,\n",
        "                                 kernel_size=1,\n",
        "                                 stride=1,\n",
        "                                 padding=0)\n",
        "\n",
        "        self.transformer_blocks = nn.ModuleList(\n",
        "            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim)\n",
        "                for d in range(depth)]\n",
        "        )\n",
        "\n",
        "        self.proj_out = zero_module(nn.Conv2d(inner_dim,\n",
        "                                              in_channels,\n",
        "                                              kernel_size=1,\n",
        "                                              stride=1,\n",
        "                                              padding=0))\n",
        "\n",
        "    def forward(self, x, context=None):\n",
        "        # note: if no context is given, cross-attention defaults to self-attention\n",
        "        b, c, h, w = x.shape\n",
        "        x_in = x\n",
        "        x = self.norm(x)\n",
        "        x = self.proj_in(x)\n",
        "        x = rearrange(x, 'b c h w -> b (h w) c')\n",
        "        for block in self.transformer_blocks:\n",
        "            x = block(x, context=context)\n",
        "        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)\n",
        "        x = self.proj_out(x)\n",
        "        return x + x_in\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "52oR7-9kcd-P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # Start stable-diffusion\n",
        "%cd /content/gdrive/MyDrive/sd/stable-diffusion/\n",
        "!python /content/gdrive/MyDrive/sd/stable-diffusion-webui/webui.py --share"
      ],
      "metadata": {
        "cellView": "form",
        "id": "PjzwxTkPSPHf"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}