fast_stable_diffusion_relaxed.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 have 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 Relaxed mode\n",
        "%%capture\n",
        "%cd /content/gdrive/MyDrive\n",
        "!git clone https://github.com/sd-webui/stable-diffusion-webui\n",
        "%cd /content/\n",
        "!git clone https://github.com/TheLastBen/fast-stable-diffusion\n",
        "!rm -R /content/gdrive/MyDrive/stable-diffusion-webui/scripts/frontend\n",
        "!cp '/content/fast-stable-diffusion/Relaxed-mode/Relaxed.py' '/content/gdrive/MyDrive/stable-diffusion-webui/scripts'\n",
        "%cd /content/gdrive/MyDrive/stable-diffusion-webui/\n",
        "!mkdir -p cache/{huggingface,torch}\n",
        "%cd /content/\n",
        "!ln -s /content/gdrive/MyDrive/stable-diffusion-webui/cache/huggingface ../root/.cache/\n",
        "!ln -s /content/gdrive/MyDrive/stable-diffusion-webui/cache/torch ../root/.cache/"
      ],
      "metadata": {
        "id": "CFWtw-6EPrKi",
        "cellView": "form"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "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/stable-diffusion-webui/models/ldm/stable-diffusion-v1/model.ckpt'\n",
        "    if os.path.exists('/content/gdrive/MyDrive/stable-diffusion-webui/models/ldm/stable-diffusion-v1/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",
        "\n",
        "  if token == \"\" and not os.path.exists('/content/gdrive/MyDrive/stable-diffusion-webui/models/ldm/stable-diffusion-v1/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/stable-diffusion-webui/models/ldm/stable-diffusion-v1/model.ckpt'\n",
        "    if os.path.exists('/content/gdrive/MyDrive/stable-diffusion-webui/models/ldm/stable-diffusion-v1/model.ckpt'):\n",
        "      clear_output()\n",
        "      print(\"Model successfully downloaded\")  \n",
        "\n",
        "  elif not os.path.exists('/content/gdrive/MyDrive/stable-diffusion-webui/models/ldm/stable-diffusion-v1/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/stable-diffusion-webui/models/ldm/stable-diffusion-v1/model.ckpt'\n",
        "        if os.path.exists('/content/gdrive/MyDrive/stable-diffusion-webui/models/ldm/stable-diffusion-v1/model.ckpt'):\n",
        "          clear_output()\n",
        "          print(\"Model successfully downloaded\")  \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": 4,
      "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/stable-diffusion-webui/src/taming-transformers/taming'):\n",
        "  %cd /content/gdrive/MyDrive/stable-diffusion-webui/\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/hlky/k-diffusion-sd#egg=k_diffusion\n",
        "  !pip install -e git+https://github.com/devilismyfriend/latent-diffusion#egg=latent-diffusion\n",
        "!pip install git+https://github.com/openai/CLIP#egg=clip\n",
        "!pip install scikit-image\n",
        "!pip install clean-fid\n",
        "!pip install accelerate\n",
        "!pip install einops\n",
        "!pip install jsonmerge\n",
        "!pip install wandb\n",
        "!pip install resize_right\n",
        "!pip install torchdiffeq\n",
        "!pip install numpy==1.21.6\n",
        "!pip install albumentations==0.4.3\n",
        "!pip install diffusers==0.3.0\n",
        "!pip install gradio\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 piexif==1.1.3\n",
        "!pip install pudb==2019.2\n",
        "!pip install pynvml==11.4.1\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 # Installing xformers\n",
        "import os\n",
        "import time\n",
        "from IPython.display import HTML\n",
        "from IPython.display import clear_output\n",
        "from subprocess import getoutput\n",
        "\n",
        "!cp -R '/content/fast-stable-diffusion/Relaxed-mode/frontend' '/content/gdrive/MyDrive/stable-diffusion-webui/scripts'\n",
        "if not os.path.exists('/content/gdrive/MyDrive/stable-diffusion-webui/scripts/k_diffusion'):\n",
        "  !cp -R '/content/gdrive/MyDrive/stable-diffusion-webui/src/taming-transformers/taming' '/content/gdrive/MyDrive/stable-diffusion-webui/scripts'\n",
        "  !cp -R '/content/gdrive/MyDrive/stable-diffusion-webui/ldm' '/content/gdrive/MyDrive/stable-diffusion-webui/scripts/ldm'\n",
        "  !cp -R '/content/gdrive/MyDrive/stable-diffusion-webui/src/k-diffusion/k_diffusion' '/content/gdrive/MyDrive/stable-diffusion-webui/scripts'\n",
        "  !cp -R '/content/gdrive/MyDrive/stable-diffusion-webui/optimizedSD' '/content/gdrive/MyDrive/stable-diffusion-webui/scripts'\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 !')"
      ],
      "metadata": {
        "id": "a---cT2rwUQj",
        "cellView": "form"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # Patching attention.py\n",
        "%%writefile /content/gdrive/MyDrive/stable-diffusion-webui/scripts/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",
        "\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",
        "\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",
        "        return self.to_out(out)\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"
      ],
      "metadata": {
        "cellView": "form",
        "id": "52oR7-9kcd-P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # Start Stable Diffusion\n",
        "%cd /content/gdrive/MyDrive/stable-diffusion-webui/\n",
        "!python /content/gdrive/MyDrive/stable-diffusion-webui/scripts/Relaxed.py"
      ],
      "metadata": {
        "cellView": "form",
        "id": "PjzwxTkPSPHf"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}