Stable_Diffusion_Helper.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/fladdict/stable-diffusion/blob/main/Stable_Diffusion_Helper.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i3Oeq07oJ28B"
      },
      "source": [
        "# Stable Diffusion Helper\n",
        "\n",
        "画像生成AI [StableDiffusion](https://github.com/CompVis/stable-diffusion)をGUIで使えるノートブックです。\n",
        "重さや初期画像などの高度機能も実装。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dBx4NIUO-B-2"
      },
      "source": [
        "## 使い方\n",
        "\n",
        "* このページ上部のメニューで、「ランタイム > ランタイムのタイプを変更」からGPUを有効化\n",
        "* [HuggingFace](https://huggingface.co/)でアカウントを作成\n",
        "* [StableDiffusionのモデルページ](https://huggingface.co/CompVis/stable-diffusion-v-1-4-original)で、「利用規約」に合意する。\n",
        "* モデルファイル [sd-v1-4.ckpt](https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt) をダウンロード\n",
        "* モデルファイルを Google Drive等にアップロード\n",
        "* 下のセル 「1-1. Google Driveとの接続」を実行\n",
        "* 下のセル　「1-2. のフォーム」に、Google Driveにアップしたモデルのパスをセット\n",
        "* このページ上部のメニューで、「ランタイム > 全てのセルを実行」を選択\n",
        "* 一番したのほうにGUIが出現する。（近くのURLで別窓でも開ける）。\n",
        "\n",
        "## 不安程な場合\n",
        "* CUDA Errorが出る場合、メモリが足りてないのでGoogle Colab Pro（or Pro+)を検討ください。\n",
        "* 一度に大量の画像生成をしてGUIが不安程な場合、GUIを実行してるセルを一時停止して、そのセルを再実行してください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kq1JsCJiuQ58"
      },
      "source": [
        "## 重要\n",
        "\n",
        "「[Stable Diffusionの利用ライセンス](https://huggingface.co/spaces/CompVis/stable-diffusion-license)」を遵守してご利用ください。\n",
        "\n",
        "----\n",
        "\n",
        "## 利用前の注意\n",
        "画像生成AIは、インターネットそのものの縮図です。あらゆるものを生成するので、生成者は自分の生成物に責任をもつ必要があります。\n",
        "\n",
        "多くの場合、問題ある画像は偶発的というよりは、「生成者が意図的に指示」をすることで生成されます。以下のようなことを心がけましょう。\n",
        "\n",
        "\n",
        "* ポルノを含む、性的な画像を生成しない（海外基準で罰せられる可能性があります）。\n",
        "* 攻撃的な画像、差別的な画像、人を不快にする目的の画像を生成しない。\n",
        "* 政治的な主張に用いない。\n",
        "* 各種の文化バイアスがかかる場合があります。生成者が適宜バランスを調整をする（例、「結婚式」の画像は欧米式で異性愛の画像になりやすい。医者の画像は白人男性になりやすい、他人の著作物をアップロードしない）。\n",
        "* 他者の権利を侵害しない（孫悟空やダースベイダーなどを意図的に作らない）。\n",
        "* 実材の人物、事件、イベントの画像（フェイクニュース含む）を作成しない。\n",
        "* 現役の作家の画風を単独指名で入力しない（個人的に推奨のマナーです）。\n",
        "\n",
        "----\n",
        "\n",
        "## お願い\n",
        "AIによる画像生成、仕事がなくなるといった文脈で煽る方向の流れは、望むものではありません。\n",
        "むしろ、みんなで「新しい創作」はどういうものか？アーティストはどうAIを使いこなしていけばいいのか？を模索していければお思います。 \n",
        "\n",
        "活版印刷が著作権の概念を生み、写真が印象派や抽象芸術の扉を開いたように、新しいテクノロジーは、新しい表現をもたらします。今、必要なことは、みんなであらゆる方向から実験をして、新しい可能性の総当たり探索をすることだと思います。\n",
        "\n",
        "そんな方向性で使ってもらえればと。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-9MvuOa_Bg9H"
      },
      "source": [
        "## 謝辞\n",
        "構成コードは以下の方々のライブラリ、スニペット、コードを参考、あるいは依拠しています。\n",
        "またnotebookは下記コード群のライセンスを継承します。\n",
        "\n",
        "* [StableDiffusion](https://github.com/CompVis/stable-diffusion) - [CreativeML Open RAIL-M License](https://github.com/CompVis/stable-diffusion/blob/main/LICENSE)\n",
        "* [Diffusers](https://github.com/huggingface/diffusers) - [Apache License 2.0](https://github.com/huggingface/diffusers/blob/main/LICENSE)\n",
        "* KLMSサンプリングは、[@RiversHaveWings](https://twitter.com/RiversHaveWings) 氏の [KLMS Sampling](https://github.com/crowsonkb/k-diffusion.git)より。 [MIT License](https://github.com/crowsonkb/k-diffusion/blob/master/LICENSE)\n",
        "* プロンプトのウェイト処理は、[@Lincoln Stein](https://github.com/lstein)氏のカスタム版[StableDiffusion](https://github.com/lstein/stable-diffusion)より。 [MIT LICENSE](https://github.com/lstein/stable-diffusion/blob/main/LICENSE)\n",
        "* KLMS連携の理解に [@pharmapsychotic](https://twitter.com/pharmapsychotic)氏の[Stable Diffusion notebook](https://colab.research.google.com/github/pharmapsychotic/ai-notebooks/blob/main/pharmapsychotic_Stable_Diffusion.ipynb#scrollTo=UU52ZvES6-1T)。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Q9zaq2YG9Gv"
      },
      "source": [
        "## 更新履歴\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e3EqPl-NQCtM"
      },
      "source": [
        "* 20220904: TextInversionで作成したembedding.ptファイルに対応\n",
        "* タイルパターンモードに対応\n",
        "* セーフフィルタのリックおじさんを空白画像に差し替え\n",
        "* プロンプトにMidJourney風ウェイト処理を追加。「dog::5 cat::3」などとできる。\n",
        "* DiffuserベースだとImg2Imgが限定的だったので、Diffuserやめる。\n",
        "* Gradioで最低限のGUIをつける"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QYSWe7iUKRc9"
      },
      "source": [
        "# セットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rcGfQNeFPP6h",
        "outputId": "0e50f4af-92c8-497d-8c16-64c6d237fe7e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "#@markdown ## 1-1. Google Driveのマウント\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "9RQ_rTSiMPpr",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@markdown ## 1-2. Google Driveにアップしたモデルのパスを設定\n",
        "#@markdown 左メニューからGoogle Driveを掘り、アップロードしたckptファイルを選択し、右クリックから「パスをコピー」を行います。\n",
        "#@markdown コピーしたパスを下のフォームにコピペしてください。\n",
        "\n",
        "GDRIVE_MODEL_PATH = \"/content/drive/MyDrive/stable-diffusion/sd-v1-4.ckpt\" #@param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "I54Hq4Ce7ddY",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@markdown ## 1-3. 画像の保存設定\n",
        "DRIVE_PATH = \"/content/drive/MyDrive\" #Driveのルート\n",
        "SAVE_FILE = True #@param {type:\"boolean\"}\n",
        "SAVE_FILE_PATH = \"/content/drive/MyDrive/stable-diffusion/output\" #@param {type:\"string\"}\n",
        "SAVE_FILE_PREFIX = \"SD\" #@param {type:\"string\"}\n",
        "\n",
        "#タイルモード記憶変数\n",
        "tile_mode_init_Conv2d = None\n",
        "tile_mode_init_ConvTranspose2d = None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uc5OwvKdjRJF",
        "outputId": "1f351de6-c7d7-454d-ce16-16302eea6193"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-ef086826-216b-491a-3d80-8bc4073ca087)\n"
          ]
        }
      ],
      "source": [
        "#GPUの確認\n",
        "!nvidia-smi -L"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w_u08HhXKZdq"
      },
      "outputs": [],
      "source": [
        "#必要ファイルのインストール\n",
        "%cd /content/\n",
        "\n",
        "#GIT\n",
        "!git clone https://github.com/CompVis/stable-diffusion  \n",
        "%cd /content/stable-diffusion\n",
        "\n",
        "!git clone https://github.com/CompVis/taming-transformers\n",
        "!git clone https://github.com/openai/CLIP.git\n",
        "!git clone https://github.com/crowsonkb/k-diffusion.git\n",
        "\n",
        "#PIP\n",
        "!pip install albumentations \n",
        "!pip install diffusers==0.2.4 \n",
        "!pip install gradio \n",
        "!pip install numpy einops kornia\n",
        "!pip install omegaconf\n",
        "!pip install pytorch-lightning\n",
        "!pip install torch-fidelity\n",
        "!pip install transformers\n",
        "!pip install ftfy jsonmerge resize-right torchdiffeq tqdm\n",
        "\n",
        "#Pathを通す\n",
        "import sys\n",
        "sys.path.append(\".\")\n",
        "sys.path.append(\"./CLIP\")\n",
        "sys.path.append('./taming-transformers')\n",
        "sys.path.append('./k-diffusion')\n",
        "sys.path.append('./ldm')\n",
        "\n",
        "#k_diffusionは初期化が必要\n",
        "!echo '' > ./k-diffusion/k_diffusion/__init__.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "DCfyEcERiwPF",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "#@markdown # タイルモードのオンオフ\n",
        "#@markdown 作成する画像をループするタイル画像にします。この設定を変えた場合、ここから下のセルを全て実行しなおす必要があります。\n",
        "\n",
        "#klassの__init__をいったん退避\n",
        "#新しい__init__を定義\n",
        "\n",
        "# code by lox9973\n",
        "# https://gitlab.com/-/snippets/2395088\n",
        "\n",
        "tile_mode = False #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "#オリジナルのConvの初期化関数を保存\n",
        "if tile_mode_init_Conv2d == None:\n",
        "  tile_mode_init_Conv2d = torch.nn.Conv2d.__init__\n",
        "  tile_mode_init_ConvTranspose2d = torch.nn.ConvTranspose2d.__init__\n",
        "\n",
        "\n",
        "def activate_tile_mode():\n",
        "  if torch.nn.Conv2d.__init__ != tile_mode_init_Conv2d:\n",
        "    return\n",
        "\n",
        "  for klass in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:\n",
        "      patch_conv(klass)\n",
        "  tile_mode = True\n",
        "  print(\"tile mode activated\")\n",
        "\n",
        "\n",
        "#Conv Filterの復元\n",
        "def deactivate_tile_mode():\n",
        "  if torch.nn.Conv2d.__init__ == tile_mode_init_Conv2d:\n",
        "    return\n",
        "  torch.nn.Conv2d.__init__ = tile_mode_init_Conv2d\n",
        "  torch.nn.ConvTranspose2d.__init__ = tile_mode_init_ConvTranspose2d\n",
        "  tile_mode == False\n",
        "  print(\"tile mode deactivated\")\n",
        "\n",
        "def patch_conv(klass):\n",
        "\tinit = klass.__init__\n",
        "\tdef __init__(self, *args, **kwargs):\n",
        "\t\treturn init(self, *args, **kwargs, padding_mode='circular')\n",
        "\tklass.__init__ = __init__\n",
        "\n",
        "if tile_mode == True:\n",
        "  activate_tile_mode()\n",
        "else:\n",
        "  deactivate_tile_mode()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bcHsbr3hblrk"
      },
      "outputs": [],
      "source": [
        "import argparse, gc, json, os, random, sys, time, glob, requests\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import numpy as np\n",
        "import PIL\n",
        "from contextlib import contextmanager, nullcontext\n",
        "from einops import rearrange, repeat\n",
        "from IPython.display import display, clear_output\n",
        "from itertools import islice\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from pytorch_lightning import seed_everything\n",
        "from torch.cuda.amp import autocast\n",
        "from ldm.util import instantiate_from_config\n",
        "from ldm.models.diffusion.ddim import DDIMSampler\n",
        "from ldm.models.diffusion.plms import PLMSSampler\n",
        "\n",
        "from k_diffusion.sampling import sample_lms\n",
        "from k_diffusion.external import CompVisDenoiser\n",
        "\n",
        "from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker\n",
        "from transformers import AutoFeatureExtractor\n",
        "\n",
        "from datetime import datetime\n",
        "import gc\n",
        "\n",
        "#メモリのクリーンアップ\n",
        "def clear_memory():\n",
        "  gc.collect()\n",
        "  torch.cuda.empty_cache() \n",
        "\n",
        "\n",
        "#新しいDenoiser\n",
        "class CFGDenoiser(nn.Module):\n",
        "    def __init__(self, model):\n",
        "        super().__init__()\n",
        "        self.inner_model = model\n",
        "\n",
        "    def forward(self, x, sigma, uncond, cond, cond_scale):\n",
        "        x_in = torch.cat([x] * 2)\n",
        "        sigma_in = torch.cat([sigma] * 2)\n",
        "        cond_in = torch.cat([uncond, cond])\n",
        "        uncond, cond = self.inner_model(x_in, sigma_in, cond=cond_in).chunk(2)\n",
        "        return uncond + (cond - uncond) * cond_scale\n",
        "\n",
        "\n",
        "#Config用クラス\n",
        "class SDOption():\n",
        "  def __init__(self):\n",
        "    self.ckpt = GDRIVE_MODEL_PATH\n",
        "    self.config = 'configs/stable-diffusion/v1-inference.yaml'\n",
        "    self.ddim_eta = 0.0\n",
        "    self.ddim_steps = 50\n",
        "    self.embedding = None # TextInversion対応用のEmbedding.pyファイルへのパス\n",
        "    self.fixed_code = True\n",
        "    self.init_img = None\n",
        "    self.init_mask = None\n",
        "    self.n_iter = 1\n",
        "    self.n_samples = 1\n",
        "    self.outdir = SAVE_FILE_PATH\n",
        "    self.precision = 'full' # 'autocast'\n",
        "    self.prompt = \"\"\n",
        "    self.sampler = 'klms'\n",
        "    self.save = True\n",
        "    self.scale = 7.5\n",
        "    self.seed = -1\n",
        "    self.strength = 0.5\n",
        "    self.variations_mode = True #Export variations of init Image\n",
        "    self.H = 512\n",
        "    self.W = 512\n",
        "    self.C = 4\n",
        "    self.f = 8\n",
        "    self.prompt_conditioning = None #promptをテキストでなくconditioningで直接注入する場合\n",
        "\n",
        "\n",
        "#第3パラメーターにEmbeddingPath（Text InversionでトレーニングしたTextEmbedderを指定可能に）\n",
        "class SDHelper():\n",
        "  def __init__(self, config_path, model_path, embedding_path = None):\n",
        "    config = OmegaConf.load(config_path)\n",
        "    self.device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "    self.model = self.load_model_from_config(config, model_path).to(self.device)\n",
        "    self.safety_model_id = \"CompVis/stable-diffusion-safety-checker\"\n",
        "    self.safety_feature_extractor = AutoFeatureExtractor.from_pretrained(self.safety_model_id)\n",
        "    self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(self.safety_model_id)\n",
        "\n",
        "    if embedding_path != None:\n",
        "      if hasattr(self.model, \"embedding_manager\"):\n",
        "        self.model.embedding_manager.load(embedding_path)\n",
        "      else:\n",
        "        print(\"Warning: Model does not have embedding_manager\")\n",
        "\n",
        "\n",
        "  def chunk(self, it, size):\n",
        "      it = iter(it)\n",
        "      return iter(lambda: tuple(islice(it, size)), ())\n",
        "\n",
        "\n",
        "  def load_model_from_config(self, config, ckpt, verbose=False):\n",
        "      print(f\"Loading model from {ckpt}\")\n",
        "      pl_sd = torch.load(ckpt, map_location=\"cpu\")\n",
        "      if \"global_step\" in pl_sd:\n",
        "          print(f\"Global Step: {pl_sd['global_step']}\")\n",
        "      sd = pl_sd[\"state_dict\"]\n",
        "      model = instantiate_from_config(config.model)\n",
        "      m, u = model.load_state_dict(sd, strict=False)\n",
        "      if len(m) > 0 and verbose:\n",
        "          print(\"missing keys:\")\n",
        "          print(m)\n",
        "      if len(u) > 0 and verbose:\n",
        "          print(\"unexpected keys:\")\n",
        "          print(u)\n",
        "\n",
        "      model.cuda()\n",
        "      model.eval()\n",
        "      return model\n",
        "\n",
        "  def make_batch(self, image, mask):\n",
        "      image = np.array(Image.open(image).convert(\"RGB\"))\n",
        "      image = image.astype(np.float32)/255.0\n",
        "      image = image[None].transpose(0,3,1,2)\n",
        "      image = torch.from_numpy(image)\n",
        "\n",
        "      mask = np.array(Image.open(mask).convert(\"L\"))\n",
        "      mask = mask.astype(np.float32)/255.0\n",
        "      mask = mask[None,None]\n",
        "      mask[mask < 0.5] = 0\n",
        "      mask[mask >= 0.5] = 1\n",
        "      mask = torch.from_numpy(mask)\n",
        "\n",
        "      masked_image = (1-mask)*image\n",
        "\n",
        "      batch = {\"image\": image, \"mask\": mask, \"masked_image\": masked_image}\n",
        "      for k in batch:\n",
        "          batch[k] = batch[k].to(device=self.device)\n",
        "          batch[k] = batch[k]*2.0-1.0\n",
        "      return batch\n",
        "\n",
        "\n",
        "  def load_img(self, path, w, h):\n",
        "    if path.startswith('http://') or path.startswith('https://'):\n",
        "        image = Image.open(requests.get(path, stream=True).raw).convert('RGB')\n",
        "    else:\n",
        "        if os.path.isdir(path):\n",
        "            files = [file for file in os.listdir(path) if file.endswith('.png') or file .endswith('.jpg')]\n",
        "            path = os.path.join(path, random.choice(files))\n",
        "            print(f\"Chose random init image {path}\")\n",
        "        image = Image.open(path).convert('RGB')\n",
        "    image = image.resize((w, h), Image.LANCZOS)\n",
        "    w, h = image.size\n",
        "    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32\n",
        "    image = image.resize((w, h), resample=PIL.Image.LANCZOS)\n",
        "    image = np.array(image).astype(np.float32) / 255.0\n",
        "    image = image[None].transpose(0, 3, 1, 2)\n",
        "    image = torch.from_numpy(image)\n",
        "    return 2.*image - 1.\n",
        "\n",
        "\n",
        "  def numpy_to_pil(self, images):\n",
        "    if images.ndim == 3:\n",
        "        images = images[None, ...]\n",
        "    images = (images * 255).round().astype(\"uint8\")\n",
        "    pil_images = [Image.fromarray(image) for image in images]\n",
        "    return pil_images\n",
        "\n",
        "\n",
        "  def load_replacement(self, x):\n",
        "      try:\n",
        "          hwc = x.shape\n",
        "          #セーフフィルターのリック・ストレイは、日本では法律に触れる可能性があるので、違う画像に差し替えます。\n",
        "          #y = Image.open(\"assets/rick.jpeg\").convert(\"RGB\").resize((hwc[1], hwc[0]))\n",
        "          y = PIL.Image.new(mode=\"RGB\", size=(hwc[1], hwc[0]))\n",
        "          y = (np.array(y)/255.0).astype(x.dtype)\n",
        "          assert y.shape == x.shape\n",
        "          return y\n",
        "      except Exception:\n",
        "          return x\n",
        "\n",
        "\n",
        "  def check_safety(self, x_image):\n",
        "    safety_checker_input = self.safety_feature_extractor(self.numpy_to_pil(x_image), return_tensors=\"pt\")\n",
        "    x_checked_image, has_nsfw_concept = self.safety_checker(images=x_image, clip_input=safety_checker_input.pixel_values)\n",
        "    assert x_checked_image.shape[0] == len(has_nsfw_concept)\n",
        "    for i in range(len(has_nsfw_concept)):\n",
        "        if has_nsfw_concept[i]:\n",
        "            x_checked_image[i] = self.load_replacement(x_checked_image[i])\n",
        "    return x_checked_image, has_nsfw_concept\n",
        "\n",
        "\n",
        "  def get_prompt_weight(self, prompt):\n",
        "    return self.model.get_learned_conditioning(prompt)\n",
        "    \n",
        "\n",
        "  def generate(self, opt):\n",
        "      global sample_idx\n",
        "      seed_everything(opt.seed)\n",
        "\n",
        "      #出力ディレクトリの作成\n",
        "      os.makedirs(opt.outdir, exist_ok=True)\n",
        "\n",
        "      #サンプラー選択\n",
        "      if opt.sampler == 'plms':\n",
        "          sampler = PLMSSampler(self.model)\n",
        "      else:\n",
        "          sampler = DDIMSampler(self.model)\n",
        "\n",
        "      model_wrap = CompVisDenoiser(self.model)       \n",
        "      batch_size = opt.n_samples\n",
        "\n",
        "      #promptをバッチの数だけコピー\n",
        "      prompt = opt.prompt\n",
        "      assert prompt is not None\n",
        "      data = [batch_size * [prompt]]\n",
        "\n",
        "\n",
        "      #初期画像の潜在空間を作成\n",
        "      init_latent = None\n",
        "      if opt.init_img != None and opt.init_img != '':\n",
        "          init_image = self.load_img(opt.init_img, opt.W, opt.H).to(self.device)\n",
        "          init_image = repeat(init_image, '1 ... -> b ...', b=batch_size)\n",
        "          init_latent = self.model.get_first_stage_encoding(self.model.encode_first_stage(init_image))  # move to latent space\n",
        "\n",
        "\n",
        "      #Inpaint用（つくりかけ）\n",
        "      \"\"\"\n",
        "      if opt.init_mask != None:\n",
        "        init_image = self.load_img(opt.init_img, opt.W, opt.H).to(self.device)\n",
        "        init_mask = self.load_img(opt.init_mask, opt.W, opt.H).to(self.device)\n",
        "        init_masked_image = (1-init_mask)*init_image\n",
        "        batch = {\"image\": init_image, \"mask\": init_mask, \"masked_image\": init_masked_image}\"\"\"\n",
        "\n",
        "      sampler.make_schedule(ddim_num_steps=opt.ddim_steps, ddim_eta=opt.ddim_eta, verbose=False)\n",
        "      t_enc = int(opt.strength * opt.ddim_steps)\n",
        "\n",
        "\n",
        "      #?? txt2imgで使われる初期値っぽいが…？\n",
        "      start_code = None\n",
        "      if opt.fixed_code and init_latent == None:\n",
        "          start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=self.device)\n",
        "\n",
        "      precision_scope = autocast if opt.precision == \"autocast\" else nullcontext\n",
        "\n",
        "      images = []\n",
        "      with torch.no_grad():\n",
        "          with precision_scope(\"cuda\"):\n",
        "              with self.model.ema_scope():\n",
        "                  for n in range(opt.n_iter):\n",
        "                      for prompts in data:\n",
        "                          #init_latentに初期画像の潜在空間\n",
        "                          #cにコンディショナル条件の潜在空間\n",
        "                          #ucにあんコンディショナルの潜在\n",
        "\n",
        "                          uc = None\n",
        "                          if opt.scale != 1.0:\n",
        "                              uc = self.model.get_learned_conditioning(batch_size * [\"\"])\n",
        "\n",
        "                          if isinstance(prompts, tuple):\n",
        "                              prompts = list(prompts)\n",
        "\n",
        "                          #プロンプトのウェイト処理\n",
        "                          #全てのプロンプトに正規化したウェイトをかけて合算する\n",
        "                          subprompts, weights = SDHelper.prompt_splitter(prompts[0])\n",
        "                          if len(subprompts) > 1:\n",
        "                            c = torch.zeros_like(uc)\n",
        "                            # get total weight for normalizing\n",
        "                            totalWeight = sum(weights)\n",
        "                            # normalize each \"sub prompt\" and add it\n",
        "                            for i in range(0,len(subprompts)):\n",
        "                              weight = weights[i]\n",
        "                              #if not skip_normalize:\n",
        "                              # skip_normalizeがついてる意図が不明なので外す。\n",
        "                              weight = weight / totalWeight\n",
        "                              c = torch.add(c,self.model.get_learned_conditioning(subprompts[i]), alpha=weight)\n",
        "                          else: # just standard prompt\n",
        "                            c = self.model.get_learned_conditioning(prompts)\n",
        "\n",
        "                          \n",
        "                          #promptをテキストでなく Tensorで直接注入する場合\n",
        "                          if (opt.prompt_conditioning != None):\n",
        "                            c = opt.prompt_conditioning\n",
        "\n",
        "\n",
        "                          if init_latent != None:\n",
        "                              #Img2Ima\n",
        "                              \n",
        "                              #z_enc に 初期画像の潜在空間\n",
        "                              #cにテキストの潜在空間\n",
        "                              #t_encに 画像のstrength * ddimsteps?\n",
        "                              #ucに空白\n",
        "                              z_enc = sampler.stochastic_encode(init_latent, torch.tensor([t_enc]*batch_size).to(self.device))\n",
        "                              samples = sampler.decode(z_enc, \n",
        "                                                       c, \n",
        "                                                       t_enc, \n",
        "                                                       unconditional_guidance_scale=opt.scale,\n",
        "                                                       unconditional_conditioning=uc)\n",
        "                          else:\n",
        "                              if opt.sampler == 'klms':\n",
        "                                  print(\"Using KLMS sampling\")\n",
        "                                  shape = [opt.C, opt.H // opt.f, opt.W // opt.f]\n",
        "                                  sigmas = model_wrap.get_sigmas(opt.ddim_steps)\n",
        "                                  model_wrap_cfg = CFGDenoiser(model_wrap)\n",
        "                                  x = torch.randn([opt.n_samples, *shape], device=self.device) * sigmas[0]\n",
        "                                  extra_args = {'cond': c, 'uncond': uc, 'cond_scale': opt.scale}\n",
        "                                  samples = sample_lms(model_wrap_cfg, \n",
        "                                                       x, \n",
        "                                                       sigmas, \n",
        "                                                       extra_args=extra_args, \n",
        "                                                       disable=False)\n",
        "                              else:\n",
        "                                  shape = [opt.C, opt.H // opt.f, opt.W // opt.f]\n",
        "                                  samples, _ = sampler.sample(S=opt.ddim_steps,\n",
        "                                                                  conditioning=c,\n",
        "                                                                  batch_size=opt.n_samples,\n",
        "                                                                  shape=shape,\n",
        "                                                                  verbose=False,\n",
        "                                                                  unconditional_guidance_scale=opt.scale,\n",
        "                                                                  unconditional_conditioning=uc,\n",
        "                                                                  eta=opt.ddim_eta,\n",
        "                                                                  x_T=start_code)\n",
        "\n",
        "                          x_samples = self.model.decode_first_stage(samples)\n",
        "                          x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)\n",
        "                          x_samples = x_samples.cpu().permute(0, 2, 3, 1).numpy()\n",
        "\n",
        "                          #Safety Checker added\n",
        "                          x_checked_image, has_nsfw_concept = self.check_safety(x_samples)\n",
        "                          x_checked_image_torch = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2)\n",
        "\n",
        "                          for x_sample in x_checked_image_torch:\n",
        "                            x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')\n",
        "                            images.append(Image.fromarray(x_sample.astype(np.uint8)))\n",
        "                            if (opt.save==True):\n",
        "                              file_id = datetime.today().strftime('%Y-%m-%d-%H-%M-%S')\n",
        "                              filepath = os.path.join(opt.outdir, f\"{SAVE_FILE_PREFIX}-{file_id}.png\")\n",
        "                              Image.fromarray(x_sample.astype(np.uint8)).save(filepath)\n",
        "                              #sample_idx += 1\n",
        "      return images\n",
        "\n",
        "  #Prompt Splitter is based on Lincoln Stein's code\n",
        "  #https://github.com/lstein/stable-diffusion/blob/main/ldm/simplet2i.py\n",
        "\n",
        "  #MidJourney互換でプロンプトの重さを処理するコード\n",
        "  #任意の文字列、 :: （スペース入るかも）（数字はいるかも）　（スペース入るかも）\n",
        "  def prompt_splitter(text):\n",
        "    \"\"\"\n",
        "    grabs all text up to the first occurrence of ':' \n",
        "    uses the grabbed text as a sub-prompt, and takes the value following ':' as weight\n",
        "    if ':' has no value defined, defaults to 1.0\n",
        "    repeats until no text remaining\n",
        "    \"\"\"\n",
        "    remaining = len(text)\n",
        "    prompts = []\n",
        "    weights = []\n",
        "    while remaining > 0:\n",
        "        if \"::\" in text:\n",
        "            idx = text.index(\"::\") # first occurrence from start\n",
        "            # grab up to index as sub-prompt\n",
        "            prompt = text[:idx]\n",
        "            remaining -= idx\n",
        "            # remove from main text\n",
        "            text = text[idx+2:]\n",
        "            # find value for weight \n",
        "            if \" \" in text:\n",
        "                idx = text.index(\" \") # first occurence\n",
        "            else: # no space, read to end\n",
        "                idx = len(text)\n",
        "            if idx != 0:\n",
        "                try:\n",
        "                    weight = float(text[:idx])\n",
        "                except: # couldn't treat as float\n",
        "                    print(f\"Warning: '{text[:idx]}' is not a value, are you missing a space?\")\n",
        "                    weight = 1.0\n",
        "            else: # no value found\n",
        "                weight = 1.0\n",
        "            # remove from main text\n",
        "            remaining -= idx\n",
        "            text = text[idx+1:]\n",
        "            # append the sub-prompt and its weight\n",
        "            prompts.append(prompt)\n",
        "            weights.append(weight)\n",
        "        else: # no : found\n",
        "            if len(text) > 0: # there is still text though\n",
        "                # take remainder as weight 1\n",
        "                prompts.append(text)\n",
        "                weights.append(1.0)\n",
        "            remaining = 0\n",
        "    print(prompts)\n",
        "    print(weights)\n",
        "    return prompts, weights\n",
        "\n",
        "\n",
        "#SDHelperのインスタンス化\n",
        "#第3引数にEmbeddingの.ptを指定すれば、TextInversionに対応可能\n",
        "opt = SDOption()\n",
        "sdh = SDHelper(opt.config, opt.ckpt, \"/content/drive/MyDrive/stable-diffusion/embeddings_gs-40000.pt\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UDqgT459ioLA"
      },
      "source": [
        "# GUIを起動"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#バッチ処理用"
      ],
      "metadata": {
        "id": "OVu85-Gi2671"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "* Init Image: ベース画像（オプション）\n",
        "* Prompt: 生成用のテキスト\n",
        "* Width: 幅\n",
        "* Height: 高さ\n",
        "* Cfg Scale: テキスト誘導の強さ（初期値 7.5)\n",
        "* Steps: 画像の描き込み時間。多いほど時間がかかり詳細になる。(初期値 50）\n",
        "* Init Image Strength: ベース画像をどれほど残すか（0: 無視 〜 1: オリジナル画像そのまま）\n",
        "* Num: 1回に生成する枚数\n",
        "* Seed: 画像の生成元となる乱数（-1にすると毎回ランダム）"
      ],
      "metadata": {
        "id": "LoIRg5p1KYR6"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w-cJQ8a1Q7ZL"
      },
      "outputs": [],
      "source": [
        "#GUIを起動\n",
        "import random\n",
        "import torch\n",
        "import gradio as gr\n",
        "\n",
        "import re\n",
        "from PIL import Image, ImageFont, ImageDraw, ImageFilter, ImageOps\n",
        "from io import BytesIO\n",
        "import base64\n",
        "import imageio\n",
        "\n",
        "clear_memory()\n",
        "\n",
        "\n",
        "def diffuse(init_image, prompt, width, height, guidance_scale, steps, init_strength, num, seed):\n",
        "  result = list()\n",
        "  \"\"\"\n",
        "  if init_image != None:\n",
        "    imageio.imwrite(\"data.png\", init_image[\"image\"])\n",
        "    imageio.imwrite(\"data_mask.png\", init_image[\"mask\"]) \n",
        "    init_image = Image.open(\"data.png\")\n",
        "    mask_image = Image.open(\"data_mask.png\")\n",
        "    #display(init_image)\n",
        "    #display(mask_image)\n",
        "  \"\"\"\n",
        "    \n",
        "  opt.init_img = init_image \n",
        "  #init_image[\"image\"]\n",
        "  #opt.init_mask = init_image[\"mask\"]\n",
        "  opt.strength = 1- init_strength\n",
        "  opt.prompt = prompt\n",
        "  opt.W = width\n",
        "  opt.H = height\n",
        "  opt.scale = guidance_scale\n",
        "  opt.ddim_steps = steps\n",
        "  opt.n_iter = 1\n",
        "  opt.save == SAVE_FILE\n",
        "  \n",
        "  for index in range(int(num)):\n",
        "    opt.seed = random.randint(0, 2**32) if seed == -1 else seed\n",
        "    image = sdh.generate(opt)[0]\n",
        "    display(image)\n",
        "    result.append(image)\n",
        "  return result\n",
        "\n",
        "\n",
        "def image_clear(image_init, strength_sli):\n",
        "  print(\"image clear\", image_init)\n",
        "\n",
        "\n",
        "def image_change(image_init, strength_sli):\n",
        "  if image_init == None:\n",
        "    return gr.Slider.update(visible=False)\n",
        "  return gr.Slider.update(visible=True)\n",
        "\n",
        "\n",
        "def set_image_to_init(images):\n",
        "  if len(images)==0:\n",
        "    return\n",
        "  try:\n",
        "    image_data = re.sub('^data:image/.+;base64,', '', images[0])\n",
        "    image = Image.open(BytesIO(base64.b64decode(image_data)))\n",
        "    return image\n",
        "  except IndexError:\n",
        "    print(\"failed to get image\")\n",
        "    return\n",
        "  \n",
        "\n",
        "#GradioによるUI起動\n",
        "with gr.Blocks() as demo:\n",
        "  with gr.Row():\n",
        "    with gr.Column():\n",
        "      with gr.Row():\n",
        "        prompt_txt = gr.Textbox(label=\"Prompt\", value=\"Beautiful detailing landscape oil painting of river and forest in the style of realism, perfect composition, golden hour\" )\n",
        "      with gr.Row():\n",
        "        with gr.Column():\n",
        "          init_img = gr.Image(value=None, source=\"upload\", interactive=True, tool=\"select\", type=\"filepath\", label=\"Init Image(option)\")\n",
        "          strength_sli = gr.Slider(label=\"Init Image Strength\", value=0.6, step=0.05, minimum=0.05, maximum=1, visible=False)\n",
        "\n",
        "      with gr.Row():\n",
        "        with gr.Column():\n",
        "          width_sli = gr.Slider(label=\"Width\", value=512, step=64, minimum=512, maximum=1024)\n",
        "          height_sli = gr.Slider(label=\"Height\", value=512, step=64, minimum=512, maximum=1024)\n",
        "          scale_sli = gr.Slider(label=\"Cfg Scale\", value=7.5, step=0.5, minimum=0, maximum=20)\n",
        "          steps_sli =gr.Slider(label=\"Steps\", value=50, step=10, minimum=10, maximum=300)\n",
        "        \n",
        "          num_num = gr.Number(label=\"Num\", value=1, minimum=1)\n",
        "          seed_num = gr.Number(label=\"Seed\", value=-1)\n",
        "          run_btn = gr.Button(\"Diffuse\", variant=\"Primary\")\n",
        "    with gr.Column():\n",
        "      gallery = gr.Gallery(elem_id=\"gallery\").style(height=\"640px\", container=True)\n",
        "      with gr.Row():\n",
        "        with gr.Column():\n",
        "          set_image_btn = gr.Button(\"Transfer to Init Image\").style(full_width=True)\n",
        "          run_btn.click(fn=diffuse, \n",
        "                        inputs=[init_img, prompt_txt, width_sli, height_sli, scale_sli, steps_sli, strength_sli, num_num, seed_num], \n",
        "                        outputs=[gallery])\n",
        "          \n",
        "          set_image_btn.click(fn=set_image_to_init, \n",
        "                    inputs=[gallery], \n",
        "                    outputs=[init_img])\n",
        "          \n",
        "          init_img.clear(fn=image_clear, inputs=[init_img, strength_sli], outputs=[strength_sli])\n",
        "          init_img.change(fn=image_change, inputs=[init_img, strength_sli], outputs=[strength_sli])\n",
        "\n",
        "demo.launch(debug=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2つのプロンプトをブレンドするサンプル\n",
        "\n",
        "犬とネコの合体動物などを作るときは、1つのプロンプトで頑張るよりも「犬の画像」と「ネコの画像」の中間ベクトルを作って画像生成するほうがよい。"
      ],
      "metadata": {
        "id": "fie9CgGCuzys"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # 実験コード\n",
        "#@markdown # 実行する場合は、skip_cell のチェックボックスを外してください\n",
        "skip_cell_0 = False #@param {type:\"boolean\"} \n",
        "\n",
        "if skip_cell_0==False:\n",
        "  #2つのPromptのブレンド\n",
        "\n",
        "  prompt0 = \"high quality landscape painting of medieval city in the style of realism, high fantasy, golden hour, haze\" #@param {type:\"string\"} \n",
        "  prompt1 = \"high quality landscape painting of medieval sacred forests in the style of realism, high fantasy, golden hour, haze\"  #@param {type:\"string\"} \n",
        "  interpolation_step = 10 #@param {type:\"integer\"}\n",
        "  seed = -1 #4324\n",
        "\n",
        "\n",
        "  opt = SDOption()\n",
        "\n",
        "  conditioning0 = sdh.get_prompt_weight(prompt0)\n",
        "  conditioning1 = sdh.get_prompt_weight(prompt1)\n",
        "\n",
        "  opt.seed = random.randint(0, 2**32) if seed == -1 else seed\n",
        "  for idx in range(interpolation_step):\n",
        "    ip = idx / interpolation_step\n",
        "    c = conditioning0 * (1-ip) + conditioning1 * ip\n",
        "    opt.prompt_conditioning = c\n",
        "    image = sdh.generate(opt)[0]\n",
        "    display(image)\n",
        "\n"
      ],
      "metadata": {
        "id": "k0lIBEVJpBXa"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "background_execution": "on",
      "collapsed_sections": [],
      "machine_shape": "hm",
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}