Kawaii_Diffusion.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/p1atdev/KawaiiDiffusion/blob/main/Kawaii_Diffusion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HeLhSShlFHmT"
      },
      "source": [
        "# Kawaii 生成機\n",
        "### 使い方\n",
        "\n",
        "上から順に再生ボタン押していく。\n",
        "ポップアップとかメッセージがでたらそれに従う。\n",
        "\n",
        "プロンプトは「指示文」ブロックのとこ。自由に変えて色々試してね。"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## セットアップ"
      ],
      "metadata": {
        "id": "PRWlvGwCjLvb"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZlXtljCCH0LW"
      },
      "outputs": [],
      "source": [
        "!bash git config --global credential.helper store"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zHkHsdtnry57"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aIrgth7sqFML"
      },
      "outputs": [],
      "source": [
        "!pip install diffusers==0.2.4\n",
        "!pip install transformers scipy ftfy\n",
        "!pip install \"ipywidgets>=7,<8\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OtrOo8YPoM2b"
      },
      "outputs": [],
      "source": [
        "from google.colab import output\n",
        "output.enable_custom_widget_manager()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8TRAh8G6sNfA"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import notebook_login\n",
        "\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xSKWBKFPArKS"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from diffusers import StableDiffusionPipeline\n",
        "\n",
        "# make sure you're logged in with `huggingface-cli login`\n",
        "pipe = StableDiffusionPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\", revision=\"fp16\", torch_dtype=torch.float16, use_auth_token=True)  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fZ1kV9mZhTXD"
      },
      "source": [
        "## セーフティーチェックを外す\n",
        "\n",
        "NSFW な写真を生成できるようになりますが、Googleアカウントを停止されても文句言わないでください。\n",
        "\n",
        "これを解除する理由としては、単に全然 NSFW でない、ピンクやペールオレンジが多く含まれた画像でも引っかかってしまい、生成の効率が悪いためです。\n",
        "\n",
        "心配ならここのブロックをコメントアウトするか削除してね。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K8ANWfD3NF5C"
      },
      "outputs": [],
      "source": [
        "def dummy(images, **kwargs): return images, False \n",
        "pipe.safety_checker = dummy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LA9myHTxbDhm"
      },
      "outputs": [],
      "source": [
        "pipe = pipe.to(\"cuda\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "by3mFc1NGc2n"
      },
      "source": [
        "ファイル保存用\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wPMif6tZEM5y",
        "outputId": "5f1849e3-06b7-449e-8a3e-b2b231404d14"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting path\n",
            "  Downloading path-16.4.0-py3-none-any.whl (26 kB)\n",
            "Installing collected packages: path\n",
            "Successfully installed path-16.4.0\n"
          ]
        }
      ],
      "source": [
        "!pip install path"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BeXjGNwSABBm"
      },
      "source": [
        "## Google Drive に接続\n",
        "接続認証のポップアップが出る"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-jdDi_uXFxW_",
        "outputId": "cbe01249-d717-49d5-c5b6-683cd97e6899"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount(\"/content/drive\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "emjLuMiR_3Yx"
      },
      "source": [
        "## 出力パス\n",
        "Google Drive での出力パス\n",
        "\n",
        "`/Diffusion/Output` は自分で好きなように変えて良い"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N6nPSxpKGCJL"
      },
      "outputs": [],
      "source": [
        "output_path = \"drive/MyDrive/Diffusion/Output\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HQFc99VGXJd5"
      },
      "source": [
        "## 画像保存\n",
        "生成方針ごとにフォルダ生成\n",
        "回数ごとにフォルダ生成"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pAxlrKvaXIq1"
      },
      "outputs": [],
      "source": [
        "from path import Path\n",
        "\n",
        "trial_title = \"Nekomimi_research\"\n",
        "trial_count = 0\n",
        "\n",
        "# 保存パス\n",
        "out_dir = Path(output_path)  / trial_title"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f3CZoA1nCqx9"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# 保存用フォルダの作成\n",
        "os.makedirs(out_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 指示文"
      ],
      "metadata": {
        "id": "B9OdbHYBhdAQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# テンプレート (あまりいじらない定型分)\n",
        "prompt_template = [\n",
        "]\n",
        "\n",
        "# 指令文\n",
        "prompt_core = [\n",
        "    \"a girl\", \n",
        "    \"standing on the beach\", \n",
        "]\n",
        "\n"
      ],
      "metadata": {
        "id": "U096XfYGhckp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## パラメータ"
      ],
      "metadata": {
        "id": "phr59ayHhmwf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "base_seed = 7186693530 # 乱数機で生成した適当な数字を入れてください。Google検索で乱数と入れれば生成できます。(0 〜 9_999_999_999 で生成)\n",
        "steps = 30 # 数字を上げると、何回も絵を微調整することになります。30前後がおすすめです。あげても必ずしも綺麗になるとは限らない。\n",
        "\n",
        "# 生成する数\n",
        "count = 10 \n",
        "\n",
        "# 画像サイズ\n",
        "width = 512\n",
        "height = 768 # 512, 768, 1024 など。 正方形以外の方が kawaii になりやすい"
      ],
      "metadata": {
        "id": "6ajdJt1shmZ_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iTgB4whDh6zW"
      },
      "source": [
        "いざ生成！"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 通常生成\n",
        "指定されたシードから順に1ずつ足していきながら指定された個数画像を生成する"
      ],
      "metadata": {
        "id": "GZWi4GELjSjp"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "COD_nH5KDjlS"
      },
      "outputs": [],
      "source": [
        "# 生成するブロおく\n",
        "from torch import autocast\n",
        "\n",
        "import pytz\n",
        "\n",
        "trial_count += 1\n",
        "\n",
        "# # 実行コード\n",
        "def run_prompt(seed, prompt_list, width, height, num_inference_steps=30):\n",
        "    prompt = \", \".join(prompt_list)\n",
        "    print(\"Generating\", seed)\n",
        "    # print(\"\\n\".join(prompt_list))\n",
        "    generator = torch.Generator(\"cuda\").manual_seed(seed)\n",
        "    with autocast(\"cuda\"):\n",
        "        image = pipe(prompt, width=width, height=height, num_inference_steps=num_inference_steps, generator=generator)[\"sample\"][0]\n",
        "    return image\n",
        "\n",
        "# プロンプト保存用\n",
        "def save_prompt(out_dir, prompt):\n",
        "    with open(out_dir / f\"{trial_count}_prompt.txt\", mode=\"w\") as fp:\n",
        "        fp.write(prompt)\n",
        "\n",
        "\n",
        "prompt_list = prompt_core[:] + prompt_template[:]\n",
        "\n",
        "# プロンプトを保存\n",
        "save_prompt(out_dir, \", \".join(prompt_list))\n",
        "\n",
        "# 繰り返し実行\n",
        "for k in range(count):\n",
        "    seed = base_seed + k\n",
        "\n",
        "    print(f\"Generate seed:{base_seed:d}\")\n",
        "\n",
        "    # 画像生成\n",
        "    image = run_prompt(seed, prompt_list, width, height, steps)\n",
        "\n",
        "    # 画像表示(邪魔ならコメントアウト)\n",
        "    display(image)\n",
        "\n",
        "    # 画像保存(保存しなければコメントアウト)\n",
        "    image.save(out_dir / f\"{trial_count}_{seed:d}.png\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qj5RmKKCjFd4"
      },
      "source": [
        "## バリエーション\n",
        "同じシード、プロンプト、異なるステップ数で生成する\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R5ObFXeXjE6m"
      },
      "outputs": [],
      "source": [
        "trial_count += 1\n",
        "out_dir = Path(output_path)  / trial_title\n",
        "\n",
        "# シードのリスト\n",
        "static_seeds = [\n",
        "  8487613867,\n",
        "  8487613869\n",
        "]\n",
        "\n",
        "# ステップのリスト(固定)\n",
        "step_list = [20, 25, 30, 32, 40, 50, 64] # 36とか入れるとなぜかエラー吐く\n",
        "\n",
        "for seed in static_seeds:\n",
        "  for setp in step_list:\n",
        "      print(f\"Variation seed:{seed:d} step:{setp:d}\")\n",
        "      image = run_prompt(seed, prompt_list, width, height, setp)\n",
        "      # display(image)\n",
        "      image.save(out_dir / f\"{trial_count}_{seed:d}_{setp:d}.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 比較検証コード\n",
        "同じ条件で、異なるワードを付与して結果を比較する"
      ],
      "metadata": {
        "id": "sLwpffNITTxA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "trial_count += 1\n",
        "out_dir = Path(output_path)  / trial_title\n",
        "\n",
        "# シードのリスト\n",
        "static_seeds = [\n",
        "    7186693533,\n",
        "    7186693538\n",
        "]\n",
        "\n",
        "# ステップのリスト(固定)\n",
        "static_step = 30\n",
        "\n",
        "# 変えるワードリスト\n",
        "words = [\n",
        "    \"\",\n",
        "    \"neko\",\n",
        "    \"cat\",\n",
        "    \"neko mimi\",\n",
        "    \"cat ear\",\n",
        "    \"catgirl\"\n",
        "]\n",
        "\n",
        "# 説明文\n",
        "def save_description(seed, step):\n",
        "  prompt_description = prompt_core[:]\n",
        "  prompt_description.append(\"[条件]\")\n",
        "  prompt_description += prompt_template[:]\n",
        "  description_text = f\"prompt: {', '.join(prompt_description)}\\nseed:{seed}\\nstep:{step}\"\n",
        "  save_prompt(out_dir, description_text)\n",
        "\n",
        "for seed in static_seeds:\n",
        "  #プロンプトの説明文の保存\n",
        "  save_description(seed, static_step)\n",
        "\n",
        "  for word in words:\n",
        "      print(f\"Difference seed:{static_step:d} step:{static_step:d} word:{word}\")\n",
        "      prompt_copy = prompt_core[:]\n",
        "      if word != \"\":\n",
        "        prompt_copy.append(word)\n",
        "      prompt_copy += prompt_template[:]\n",
        "\n",
        "      print(prompt_copy)\n",
        "\n",
        "      image = run_prompt(seed, prompt_copy, width, height, static_step)\n",
        "\n",
        "      # display(image)\n",
        "      image.save(out_dir / f\"{trial_count}_{seed:d}_{static_step:d}_{word}.png\")"
      ],
      "metadata": {
        "id": "0F6VRSo1TbRX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5RlqXbLaamIS"
      },
      "source": [
        "## ユーティリティ\n",
        "\n",
        "文章をカンマで区切って出力してくれるやつ。\n",
        "Lexicaのプロンプトをコピペしたい時用。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eeR5Wy55aliy"
      },
      "outputs": [],
      "source": [
        "raw_text = \"\"\"\n",
        "\"\"\"\n",
        "\n",
        "result = raw_text.strip().split(\",\")\n",
        "\n",
        "print(\"[\")\n",
        "\n",
        "for text in result:\n",
        "  print(f'\"{text.strip()}\",')\n",
        "\n",
        "print(\"]\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "PRWlvGwCjLvb",
        "fZ1kV9mZhTXD",
        "BeXjGNwSABBm",
        "emjLuMiR_3Yx",
        "HQFc99VGXJd5",
        "5RlqXbLaamIS"
      ],
      "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
}